"""한전 공문 검토·작성 API (FastAPI + onnxruntime)

환경 변수:
  MODEL_REPO  — HF 모델 repo (기본: onnx-community/EXAONE-3.5-2.4B-Instruct)
"""
import json
import os
import threading
import time
from pathlib import Path
from typing import Optional

import re
import tempfile
import xml.etree.ElementTree as ET
import zipfile

from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

MODEL_REPO = os.getenv("MODEL_REPO", "onnx-community/EXAONE-3.5-2.4B-Instruct")
DATAS_DIR = Path(__file__).parent / "datas"
DOC_TYPES = {"공문-외부발송", "공문-내부", "보고서"}

np = None
ort = None

NUM_LAYERS = 30
NUM_KV_HEADS = 8
HEAD_DIM = 80

app = FastAPI(title="공문 검토·작성 API", version="0.2")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


class Generator:
    def __init__(self):
        self.tokenizer = None
        self.session = None
        self.output_index = {}
        self.model_dir = None
        self.ready = False
        self.status = "initializing"

    def load(self):
        import numpy as _np
        import onnxruntime as _ort
        from huggingface_hub import snapshot_download
        from transformers import AutoTokenizer
        global np, ort
        np = _np
        ort = _ort

        try:
            self.status = "downloading model"
            self.model_dir = Path(snapshot_download(
                MODEL_REPO,
                allow_patterns=[
                    "config.json",
                    "generation_config.json",
                    "tokenizer.json",
                    "tokenizer_config.json",
                    "special_tokens_map.json",
                    "onnx/model_q4.onnx",
                    "onnx/model_q4.onnx_data",
                ],
            ))
            self.status = "loading tokenizer"
            self.tokenizer = AutoTokenizer.from_pretrained(
                str(self.model_dir), trust_remote_code=True
            )

            self.status = "creating ONNX session"
            so = ort.SessionOptions()
            so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
            try:
                actual_cpus = len(os.sched_getaffinity(0))
            except AttributeError:
                actual_cpus = os.cpu_count() or 2
            n_threads = int(os.getenv("OMP_NUM_THREADS", str(actual_cpus)))
            so.intra_op_num_threads = max(1, n_threads)
            so.inter_op_num_threads = 1
            print(f"[ONNX] threads: intra={so.intra_op_num_threads}, inter=1 (detected {actual_cpus} CPUs)")
            self.session = ort.InferenceSession(
                str(self.model_dir / "onnx" / "model_q4.onnx"),
                sess_options=so,
                providers=["CPUExecutionProvider"],
            )
            self.output_index = {
                o.name: i for i, o in enumerate(self.session.get_outputs())
            }
            self.ready = True
            self.status = "ready"
        except Exception as e:
            self.status = f"failed: {e}"
            raise

    def _sample(self, logits, temperature, top_p, top_k):
        logits = logits.astype(np.float32)
        if temperature <= 0:
            return int(np.argmax(logits))
        logits = logits / temperature
        if 0 < top_k < logits.size:
            kth = np.partition(logits, -top_k)[-top_k]
            logits = np.where(logits < kth, -np.inf, logits)
        probs = np.exp(logits - logits.max())
        probs = probs / probs.sum()
        if top_p < 1.0:
            order = np.argsort(-probs)
            sorted_p = probs[order]
            cum = np.cumsum(sorted_p)
            cutoff = np.searchsorted(cum, top_p) + 1
            keep_idx = order[:cutoff]
            new_probs = np.zeros_like(probs)
            new_probs[keep_idx] = probs[keep_idx]
            probs = new_probs / new_probs.sum()
        return int(np.random.choice(len(probs), p=probs))

    def _format_exaone(self, messages):
        parts = []
        for m in messages:
            role = m["role"]
            content = m["content"]
            if role == "user":
                parts.append(f"[|{role}|]{content}\n")
            else:
                parts.append(f"[|{role}|]{content}[|endofturn|]\n")
        parts.append("[|assistant|]")
        return "".join(parts)

    def stream(self, messages, max_new_tokens, temperature, top_p, top_k):
        tok = self.tokenizer
        try:
            ids = tok.apply_chat_template(
                messages, return_tensors="np", add_generation_prompt=True
            ).astype(np.int64)
        except Exception as e:
            print(f"[WARN] apply_chat_template failed ({e}), using manual format")
            prompt_str = self._format_exaone(messages)
            ids = tok(prompt_str, return_tensors="np").input_ids.astype(np.int64)
        input_ids = ids
        seq_len = input_ids.shape[1]
        attention_mask = np.ones((1, seq_len), dtype=np.int64)
        position_ids = np.arange(seq_len, dtype=np.int64).reshape(1, -1)

        past_kv = {}
        for i in range(NUM_LAYERS):
            past_kv[f"past_key_values.{i}.key"] = np.zeros(
                (1, NUM_KV_HEADS, 0, HEAD_DIM), dtype=np.float32
            )
            past_kv[f"past_key_values.{i}.value"] = np.zeros(
                (1, NUM_KV_HEADS, 0, HEAD_DIM), dtype=np.float32
            )

        eos_id = tok.eos_token_id
        total_len = seq_len

        for step in range(max_new_tokens):
            feeds = {
                "input_ids": input_ids,
                "attention_mask": attention_mask,
                "position_ids": position_ids,
                **past_kv,
            }
            outputs = self.session.run(None, feeds)
            logits = outputs[self.output_index["logits"]]
            next_id = self._sample(logits[0, -1, :], temperature, top_p, top_k)

            if next_id == eos_id:
                break

            piece = tok.decode([next_id], skip_special_tokens=True)
            yield piece

            input_ids = np.array([[next_id]], dtype=np.int64)
            total_len += 1
            attention_mask = np.ones((1, total_len), dtype=np.int64)
            position_ids = np.array([[total_len - 1]], dtype=np.int64)
            past_kv = {}
            for i in range(NUM_LAYERS):
                past_kv[f"past_key_values.{i}.key"] = outputs[
                    self.output_index[f"present.{i}.key"]
                ]
                past_kv[f"past_key_values.{i}.value"] = outputs[
                    self.output_index[f"present.{i}.value"]
                ]


# ============================================================
# 글로벌 인스턴스 + 락
# ============================================================
generator = Generator()
gen_lock = threading.Lock()


@app.on_event("startup")
def _startup():
    threading.Thread(target=generator.load, daemon=True).start()


# ============================================================
# API 스키마
# ============================================================
class Message(BaseModel):
    role: str
    content: str


class GenerateRequest(BaseModel):
    messages: list[Message]
    max_new_tokens: int = 600
    temperature: float = 0.7
    top_p: float = 0.8
    top_k: int = 20


class HwpxRequest(BaseModel):
    text: str


# ============================================================
# 엔드포인트
# ============================================================
STATIC_DIR = Path(__file__).parent / "static"
if STATIC_DIR.exists():
    app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


@app.get("/")
def root():
    index = STATIC_DIR / "index.html"
    if index.exists():
        return FileResponse(str(index))
    return {
        "status": generator.status,
        "ready": generator.ready,
        "model_repo": MODEL_REPO,
    }


@app.get("/status")
def status():
    return {
        "status": generator.status,
        "ready": generator.ready,
        "model_repo": MODEL_REPO,
    }


@app.get("/health")
def health():
    return {"ready": generator.ready, "status": generator.status}


# ============================================================
# 파일 텍스트 추출 헬퍼
# ============================================================
def _ext_hwpx(path: Path) -> str:
    HP = "http://www.hancom.co.kr/hwpml/2012/paragraph"
    lines = []
    with zipfile.ZipFile(str(path), "r") as zf:
        sections = sorted(
            n for n in zf.namelist()
            if n.startswith("Contents/section") and n.endswith(".xml")
        )
        for name in sections:
            with zf.open(name) as f:
                try:
                    root = ET.parse(f).getroot()
                except ET.ParseError:
                    continue
                for para in root.iter(f"{{{HP}}}p"):
                    lines.append("".join(
                        t.text for t in para.iter(f"{{{HP}}}t") if t.text
                    ))
    text = "\n".join(lines)
    return re.sub(r"\n{3,}", "\n\n", text).strip()


def _ext_hwp(path: Path) -> str:
    import struct, zlib
    try:
        import olefile
    except ImportError:
        raise RuntimeError("olefile 미설치 (pip install olefile)")

    try:
        ole = olefile.OleFileIO(str(path))
    except Exception as exc:
        raise RuntimeError(f"HWP OLE 열기 실패 [{type(exc).__name__}]: {exc or 'no detail'}")

    compressed = True
    try:
        if ole.exists("FileHeader"):
            hdr = ole.openstream("FileHeader").read()
            if len(hdr) >= 40:
                compressed = bool(struct.unpack_from("<I", hdr, 36)[0] & 0x1)
    except Exception:
        pass

    # Discover sections: case-insensitive, any naming convention
    try:
        all_entries = ole.listdir()
        print(f"[HWP] OLE entries: {all_entries[:20]}", flush=True)
        sections = sorted(
            (e[1] for e in all_entries
             if isinstance(e, (list, tuple)) and len(e) >= 2
             and str(e[0]).lower() == "bodytext"
             and str(e[1]).lower().startswith("section")),
            key=lambda s: int(re.sub(r"\D", "", s) or "0"),
        )
    except Exception as exc:
        raise RuntimeError(f"섹션 목록 오류 [{type(exc).__name__}]: {exc or 'no detail'}")

    if not sections:
        entry_dump = str([str(e) for e in all_entries[:20]])
        raise RuntimeError(f"BodyText/Section* 없음. OLE 항목: {entry_dump}")

    lines: list[str] = []
    for sec_name in sections:
        try:
            raw = ole.openstream(f"BodyText/{sec_name}").read()
        except Exception as exc:
            print(f"[HWP] openstream {sec_name} 오류: {exc}", flush=True)
            continue

        try:
            data = zlib.decompress(raw, -15) if compressed else raw
        except zlib.error:
            try:
                data = zlib.decompress(raw)
            except zlib.error:
                data = raw

        pos = 0
        while pos + 4 <= len(data):
            hval = struct.unpack_from("<I", data, pos)[0]
            tag_id = hval & 0x3FF
            size   = (hval >> 20) & 0xFFF
            pos += 4
            if size == 0xFFF:
                if pos + 4 > len(data): break
                size = struct.unpack_from("<I", data, pos)[0]
                pos += 4

            if tag_id == 67 and size >= 2:  # HWPTAG_PARA_TEXT
                chunk = data[pos:pos + size]
                t, chars = 0, []
                while t + 2 <= len(chunk):
                    cp = struct.unpack_from("<H", chunk, t)[0]
                    t += 2
                    if cp == 0x0D:
                        lines.append("".join(chars)); chars = []
                    elif 0x01 <= cp <= 0x0C:
                        t += 8
                    elif cp >= 0x20:
                        chars.append(chr(cp))
                if chars:
                    lines.append("".join(chars))
            pos += size

    return re.sub(r"\n{3,}", "\n\n", "\n".join(l for l in lines if l)).strip()


def _ext_docx(path: Path) -> str:
    from docx import Document
    doc = Document(str(path))
    return "\n".join(p.text for p in doc.paragraphs)


def _ext_pdf(path: Path) -> str:
    from pypdf import PdfReader
    reader = PdfReader(str(path))
    pages = [f"--- {i+1}페이지 ---\n{p.extract_text() or ''}"
             for i, p in enumerate(reader.pages)]
    return "\n".join(pages)


def _ext_txt(path: Path) -> str:
    for enc in ["utf-8", "utf-8-sig", "cp949", "euc-kr"]:
        try:
            return path.read_text(encoding=enc)
        except UnicodeDecodeError:
            continue
    raise RuntimeError("인코딩 감지 실패")


def _load_examples(doc_type: str, max_count: int = 3, max_chars: int = 2000) -> str:
    folder = DATAS_DIR / doc_type
    if not folder.exists():
        return ""
    examples = []
    for p in sorted(folder.iterdir()):
        if p.name.startswith(".") or p.stat().st_size == 0:
            continue
        try:
            ext = p.suffix.lower()
            if ext == ".hwpx":   text = _ext_hwpx(p)
            elif ext == ".hwp":  text = _ext_hwp(p)
            elif ext == ".docx": text = _ext_docx(p)
            elif ext == ".pdf":  text = _ext_pdf(p)
            elif ext in (".txt", ".md"): text = _ext_txt(p)
            else: continue
            text = text.strip()
            if text:
                examples.append(f"--- 예시: {p.name} ---\n{text[:max_chars]}")
        except Exception:
            continue
        if len(examples) >= max_count:
            break
    return "\n\n".join(examples)


@app.get("/examples/{doc_type}")
def get_examples(doc_type: str):
    if doc_type not in DOC_TYPES:
        raise HTTPException(400, f"알 수 없는 문서 유형: {doc_type}")
    return {"doc_type": doc_type, "examples": _load_examples(doc_type)}


@app.post("/extract")
async def extract_file(file: UploadFile = File(...)):
    """업로드된 문서(.hwp/.hwpx/.docx/.pdf/.txt)에서 텍스트를 추출해 반환."""
    ext = Path(file.filename or "").suffix.lower()
    supported = {".hwp", ".hwpx", ".docx", ".pdf", ".txt"}
    if ext not in supported:
        raise HTTPException(400, f"지원하지 않는 형식: {ext}. 지원: {', '.join(sorted(supported))}")

    content = await file.read()
    with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp:
        tmp.write(content)
        tmp_path = Path(tmp.name)

    try:
        if ext == ".hwpx":
            text = _ext_hwpx(tmp_path)
        elif ext == ".hwp":
            text = _ext_hwp(tmp_path)
        elif ext == ".docx":
            text = _ext_docx(tmp_path)
        elif ext == ".pdf":
            text = _ext_pdf(tmp_path)
        else:
            text = _ext_txt(tmp_path)
    except Exception as e:
        detail = str(e).strip() or type(e).__name__
        raise HTTPException(500, detail)
    finally:
        tmp_path.unlink(missing_ok=True)

    return {"text": text, "filename": file.filename}


@app.post("/generate")
def generate(req: GenerateRequest):
    """SSE 스트리밍 — 토큰별로 'data: {"token": "..."}\\n\\n' 전송."""
    if not generator.ready:
        raise HTTPException(503, f"model not ready ({generator.status})")

    if not gen_lock.acquire(blocking=False):
        raise HTTPException(429, "another generation in progress (CPU 1요청 동시 처리)")

    def event_stream():
        try:
            t0 = time.time()
            count = 0
            for piece in generator.stream(
                [m.dict() for m in req.messages],
                req.max_new_tokens,
                req.temperature,
                req.top_p,
                req.top_k,
            ):
                count += 1
                yield f"data: {json.dumps({'token': piece}, ensure_ascii=False)}\n\n"
            elapsed = time.time() - t0
            yield f"data: {json.dumps({'done': True, 'tokens': count, 'elapsed': round(elapsed, 1)}, ensure_ascii=False)}\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'error': str(e)}, ensure_ascii=False)}\n\n"
        finally:
            gen_lock.release()

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@app.post("/download/hwpx")
def download_hwpx(req: HwpxRequest):
    """텍스트를 .hwpx 파일로 다운로드."""
    import re
    from urllib.parse import quote
    from fastapi.responses import Response

    m = re.search(r"^제\s*목\s+(.+)$", req.text, re.MULTILINE)
    title = m.group(1).strip() if m else "문서"
    safe = re.sub(r'[\\/:*?"<>|]', "_", title)[:50]
    encoded = quote(safe, safe="")

    content = req.text.encode("utf-8-sig")

    return Response(
        content=content,
        media_type="application/octet-stream",
        headers={
            "Content-Disposition": f"attachment; filename*=UTF-8''{encoded}.hwpx",
        },
    )