# WygLore Leaf — Korean NER endpoint (HF Docker Space)
#
# Gradio 테스트 UI (/) + 깔끔한 FastAPI JSON API (/extract · /extract_batch · /health · /docs).
# 플러그인은 /extract 를 host.nativeFetch 로 한 방 호출 (모바일 저메모리 폴백).
#
# 프라이버시 (NSFW-민감 유저 대상 — 코드로 보장):
#   - 사용자 텍스트를 *로그/저장/학습에 일절 안 씀*. print 는 모델 메타뿐. 추론은 stateless.
#   - gradio analytics OFF · uvicorn access_log OFF (요청 IP 도 안 남김) · HF Space HTTPS.
#   - 단 텍스트가 HF 인프라를 *경유*는 함 -> 서버 경로는 플러그인에서 opt-in + 고지 (자동 X).

import gradio as gr
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from huggingface_hub import list_repo_files
from optimum.onnxruntime import ORTModelForTokenClassification
from transformers import AutoTokenizer, pipeline
from span_postproc import trim_span_offsets, is_stopword, nfc, has_hangul

MODEL_REPO = "m6dd8m/wl-ko-ner-v5"                        # ONNX 가중치 (v5 멀티봇 + Q3 후처리)
TOK_REPO   = "monologg/koelectra-base-v3-discriminator"  # 토크나이저 (파인튜닝해도 동일)

# 1) repo 안의 .onnx 자동 탐지 — fp32 우선.
onnx_files = [f for f in list_repo_files(MODEL_REPO) if f.endswith(".onnx")]
if not onnx_files:
    raise RuntimeError(f"{MODEL_REPO} 에 .onnx 가 없음 — repo 내용 확인")
onnx_path = next((f for f in onnx_files if "fp32" in f.lower()), onnx_files[0])
subfolder, onnx_name = (onnx_path.rsplit("/", 1) if "/" in onnx_path else ("", onnx_path))
print(f"[boot] onnx -> '{onnx_path}'")   # 모델 메타만; 사용자 텍스트 절대 로그 X

# 2) 모델 + 토크나이저 + pipeline.
model = ORTModelForTokenClassification.from_pretrained(
    MODEL_REPO, file_name=onnx_name, subfolder=subfolder)
tok = AutoTokenizer.from_pretrained(TOK_REPO)
print(f"[boot] id2label = {model.config.id2label}")
ner = pipeline("token-classification", model=model, tokenizer=tok,
               aggregation_strategy="simple")

# 3) pipeline 출력 -> on-device ner.js 와 동일 모양 + Q3 CHEAP 후처리 (span_postproc 공유 로직).
_TRIM = ("PS", "OG", "LC")   # 한자 주석/괄호 꼬리 트리밍은 고유명에만 (DT/TI/QT 는 괄호 안 숫자가 정당)

def _spans(entities, text):
    out = []
    for e in entities:
        typ = e["entity_group"]                              # PS / LC / OG / DT / TI / QT
        s, en = int(e["start"]), int(e["end"])
        if typ in _TRIM:
            s, en, surf = trim_span_offsets(text, s, en)     # CHEAP-1 한자괄호 트리밍 (surface+offset)
            if not has_hangul(surf):                          # 한글 0 = 한자 주석 잔재(pipeline 분리) → 드롭
                continue
        else:
            surf = e["word"]
        surf = (surf or "").strip()
        if not surf or is_stopword(surf, typ):               # CHEAP-3 stoplist + 빈 surface 드롭
            continue
        out.append({"text":  surf,
                    "type":  typ,
                    "score": float(e["score"]),
                    "start": s,
                    "end":   en})
    return out

def extract(text):
    if not text or not text.strip():
        return []
    text = nfc(text)                                         # CHEAP-4 NFC 통일 (offset 기준 일관)
    return _spans(ner(text), text)

# ───────────────────────── FastAPI (깔끔한 JSON API) ─────────────────────────
app = FastAPI(title="WygLore Leaf Korean NER")
app.add_middleware(CORSMiddleware, allow_origins=["*"],
                   allow_methods=["*"], allow_headers=["*"])

class OneReq(BaseModel):
    text: str

class BatchReq(BaseModel):
    texts: list[str]

@app.get("/health")
def health():
    return {"ok": True, "model": MODEL_REPO, "onnx": onnx_path}

@app.post("/extract")            # 한 방: {"text": "..."} -> {"entities": [...]}
def extract_api(req: OneReq):
    return {"entities": extract(req.text)}

@app.post("/extract_batch")      # 배치: {"texts": [...]} -> {"results": [[...], ...]}
def extract_batch_api(req: BatchReq):
    if len(req.texts) > 512:     # 가벼운 abuse/메모리 가드
        raise HTTPException(400, "batch too large (max 512)")
    texts = [nfc(t) if (t and t.strip()) else t for t in req.texts]   # CHEAP-4
    idx = [i for i, t in enumerate(texts) if t and t.strip()]
    outs = ner([texts[i] for i in idx]) if idx else []
    results = [[] for _ in req.texts]
    for k, i in enumerate(idx):
        results[i] = _spans(outs[k], texts[i])
    return {"results": results}

# ───────────────────────── Gradio 테스트 UI 를 / 에 마운트 ─────────────────────────
demo = gr.Interface(
    fn=extract,
    inputs=gr.Textbox(lines=4, label="한국어 텍스트",
                      placeholder="새봄이 여의도에서 안도현을 만났다."),
    outputs=gr.JSON(label="엔티티 스팬"),
    title="WygLore Leaf — Korean NER (wl-ko-ner-v2)",
    description="UI=테스트용 · API=POST /extract · 배치=/extract_batch · 문서=/docs",
    analytics_enabled=False,   # gradio 텔레메트리 OFF (프라이버시)
)
app = gr.mount_gradio_app(app, demo, path="/")  # API 라우트가 먼저 등록돼 우선 매칭

if __name__ == "__main__":
    import uvicorn
    # Docker Space: SSR 프록시 없음 -> 7860 깨끗이 바인드. access_log OFF (IP 안 남김).
    uvicorn.run(app, host="0.0.0.0", port=7860, access_log=False)