"""
app.py - FastAPI server: serves the UI at / and the OCR API at /api/*.

    uvicorn app:app --host 0.0.0.0 --port 8000          # real model (GPU)
    OCR_MOCK=1 uvicorn app:app --port 8000               # UI development, no GPU
"""
import difflib, os, re, shutil, tempfile, threading, time, uuid
from pathlib import Path

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

from ocr_engine import CONFIG, IMAGE_EXT, LANG_NAMES, OCREngine, mask_ids

try:                                   # benchmark.py needs validate.py; without it, no ground-truth scoring
    from benchmark import score
except ImportError:
    score = None

MAX_UPLOAD_MB = int(os.getenv("OCR_MAX_UPLOAD_MB", "25"))
JOB_TTL_S = int(os.getenv("OCR_JOB_TTL_S", "3600"))      # uploads + outputs deleted after this
JOBS_DIR = Path(os.getenv("OCR_JOBS_DIR", tempfile.gettempdir())) / "ocr_jobs"
JOBS_DIR.mkdir(parents=True, exist_ok=True)
STATIC = Path(__file__).parent / "static"

app = FastAPI(title="Indic Handwritten OCR")
ENGINE: OCREngine | None = None
_load_error: str | None = None


def get_engine() -> OCREngine:
    """Loads the model once. Modal / Colab can also inject a pre-loaded engine via set_engine()."""
    global ENGINE, _load_error
    if ENGINE is None:
        try:
            ENGINE = OCREngine.load()
        except Exception as e:                      # surface load failures to /api/health
            _load_error = repr(e)
            raise
    return ENGINE


def set_engine(engine: OCREngine):
    global ENGINE
    ENGINE = engine


@app.on_event("startup")
def _warm():
    if os.getenv("OCR_EAGER_LOAD", "1") == "1":
        threading.Thread(target=get_engine, daemon=True).start()
    threading.Thread(target=_janitor, daemon=True).start()


def _janitor():
    while True:
        cutoff = time.time() - JOB_TTL_S
        for d in JOBS_DIR.iterdir():
            if d.is_dir() and d.stat().st_mtime < cutoff:
                shutil.rmtree(d, ignore_errors=True)
        time.sleep(300)


@app.get("/api/health")
def health():
    return {"ready": bool(ENGINE and ENGINE.ready), "mock": CONFIG["MOCK"], "error": _load_error,
            "languages": LANG_NAMES, "max_upload_mb": MAX_UPLOAD_MB, "max_pages": CONFIG["MAX_PAGES"],
            "scoring": score is not None}


def evaluate(result: dict, gt: str, lang: str | None) -> dict:
    """Same metrics as benchmark.py's run_docs(), plus a word diff for the UI."""
    pred = "\n".join(b["text"] for p in result["pages"] for b in p["blocks"])
    gt = mask_ids(gt.strip())                      # the output is masked, so mask the ground truth the same way
    tok = lambda s: re.findall(r"\S+", s)
    g, p = tok(gt), tok(pred)
    diff = []
    for op, i0, i1, j0, j1 in difflib.SequenceMatcher(None, g, p, autojunk=False).get_opcodes():
        if op == "equal":
            diff.append(["eq", " ".join(g[i0:i1])])
        if op in ("delete", "replace"):
            diff.append(["del", " ".join(g[i0:i1])])   # in ground truth, missing from OCR
        if op in ("insert", "replace"):
            diff.append(["ins", " ".join(p[j0:j1])])   # in OCR, not in ground truth
    return {**score(gt, pred, lang or "-"), "gt_chars": len(gt), "gt_words": len(g), "diff": diff}


@app.post("/api/ocr")
def ocr(file: UploadFile = File(...), lang: str = Form(""), gt: UploadFile | None = File(None)):
    # sync def -> FastAPI runs it in a worker thread; the engine lock serialises GPU access
    suffix = Path(file.filename or "").suffix.lower()
    if suffix not in IMAGE_EXT | {".pdf"}:
        raise HTTPException(400, f"Unsupported file type '{suffix}'. Use PDF or an image.")
    job_id = uuid.uuid4().hex[:12]
    job_dir = JOBS_DIR / job_id
    job_dir.mkdir()
    safe_name = Path(file.filename).name.replace(" ", "_")
    src = job_dir / safe_name
    size = 0
    with src.open("wb") as fh:
        while chunk := file.file.read(1 << 20):
            size += len(chunk)
            if size > MAX_UPLOAD_MB << 20:
                shutil.rmtree(job_dir, ignore_errors=True)
                raise HTTPException(413, f"File is larger than {MAX_UPLOAD_MB} MB")
            fh.write(chunk)
    gt_text = None
    if gt is not None and gt.filename:
        if score is None:
            shutil.rmtree(job_dir, ignore_errors=True)
            raise HTTPException(400, "Ground-truth scoring needs validate.py and benchmark.py on the server")
        try:
            gt_text = gt.file.read(5 << 20).decode("utf-8-sig")
        except UnicodeDecodeError:
            shutil.rmtree(job_dir, ignore_errors=True)
            raise HTTPException(400, "Ground truth must be a UTF-8 .txt file")
    try:
        result = get_engine().process(src, lang or None, job_dir / "out")
    except ValueError as e:
        shutil.rmtree(job_dir, ignore_errors=True)
        raise HTTPException(400, str(e))
    except Exception as e:
        shutil.rmtree(job_dir, ignore_errors=True)
        raise HTTPException(500, f"OCR failed: {e}")
    src.unlink(missing_ok=True)                    # don't keep the uploaded document around
    result["job_id"] = job_id
    if gt_text is not None:
        result["eval"] = evaluate(result, gt_text, lang or None)
    return JSONResponse(result)


@app.get("/api/jobs/{job_id}/{name}")
def download(job_id: str, name: str):
    if not job_id.isalnum() or "/" in name or ".." in name:
        raise HTTPException(400, "bad path")
    p = JOBS_DIR / job_id / "out" / name
    if not p.is_file():
        raise HTTPException(404, "expired or not found")
    return FileResponse(p, filename=name)


app.mount("/", StaticFiles(directory=STATIC, html=True), name="ui")