app.py

"""HTR VLM + MoE post-correction — FastAPI backend.

Two modes:
  - OCR: pick one VLM, transcribe a page, correct it inline, accumulate ICL
  - Post-correction: image + OCR text → 2 expert VLMs + 1 judge VLM consensus

State is in-memory, single-process (HF Spaces single-user demo).
"""
from __future__ import annotations

import json
import os
import uuid
from typing import Any, Optional

from fastapi import FastAPI, File, Header, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse, PlainTextResponse
from fastapi.staticfiles import StaticFiles

from io_utils import (
    compute_cer_wer, diff_words, export_icl_jsonl,
    load_image_bytes_to_jpeg_b64, parse_lines_from_model_response,
)
from moe import run_moe_correction
from paths import STATIC_DIR, DATA_DIR
from prompts import (
    ICLPool, OCR_SYSTEM_TPL, OCR_USER_ZERO_SHOT_TPL, EXPERT_SYSTEM_TPL,
    EXPERT_USER_TPL, JUDGE_SYSTEM_TPL, JUDGE_USER_TPL,
    render_ocr_few_shot_header, render_ocr_system, render_ocr_user_target,
    list_preset_names, get_preset,
)
from provider import LLMClient, image_part, test_connection_sync, text_part
from schemas import (
    CURATED_MODELS, AnnotationLabelReq, CorrectionReq, OcrReq, PasteIclReq, PostCorrReq,
    SamPointReq, SettingsReq, TestKeyReq, DEFAULT_EXPERT_A, DEFAULT_EXPERT_B,
    DEFAULT_GUIDELINES, DEFAULT_JSON_TEMPLATE, DEFAULT_JUDGE, DEFAULT_LANGUAGE,
    DEFAULT_OCR_MODEL, DEFAULT_OUTPUT_MODE,
)

ENV_OPENROUTER_KEY = os.environ.get("OPENROUTER_API_KEY", "")


def _default_settings() -> dict:
    return {
        "ocr_model": DEFAULT_OCR_MODEL,
        "moe_expert_a": DEFAULT_EXPERT_A,
        "moe_expert_b": DEFAULT_EXPERT_B,
        "moe_judge": DEFAULT_JUDGE,
        "language": DEFAULT_LANGUAGE,
        "guidelines": DEFAULT_GUIDELINES,
        "output_mode": DEFAULT_OUTPUT_MODE,
        "output_json_template": DEFAULT_JSON_TEMPLATE,
        "temperature": 0.0,
        "n_icl": 2,
        "use_icl": True,
        "custom_ocr_system": None,
        "custom_ocr_user": None,
        "custom_expert_system": None,
        "custom_expert_user": None,
        "custom_judge_system": None,
        "custom_judge_user": None,
    }


SESSION: dict[str, Any] = {
    "settings": _default_settings(),
    "images": {},     # image_id -> {"b64": str, "filename": str, "w": int, "h": int}
    "pages": [],      # list of dicts (see _new_page())
    "icl_pool": ICLPool(),
    "mode": "ocr",
    "totals": {"requests": 0, "input_tokens": 0, "output_tokens": 0, "cost_usd": 0.0},
}


# Per-sandbox metadata. `language` is the per-page tag — the global ICL pool
# is filtered by exact language match, so a sandbox naturally picks up
# only its own corpus examples (Baybars 0-shot, Rasam-2 2-shot, etc.).
SANDBOX_SPECS = [
    {
        "file": "rasam2-BULAC_MS_ARA_6_49817.jpg",
        "short": "Arabic",
        "title": "RASAM-2 (Arabic manuscript, 2 shots)",
        "language": "Arabic — Rasam corpus",
        "citation_url": "https://huggingface.co/datasets/calfa-ai/RASAM-2",
    },
    {
        "file": "endp-FRAN_0393_02960_L.jpg",
        "short": "French",
        "title": "e-Notre-Dame de Paris (French manuscript, 2 shots)",
        "language": "Middle French — E-NDP",
        "citation_url": "https://zenodo.org/records/7575693",
    },
    {
        "file": "chiknowpo-BULAC_BIULO_CHI_1087_1_0409.jpg",
        "short": "Chinese",
        "title": "ChiKnowPo (Chinese xylograph, 1 shot)",
        "language": "Classical Chinese — ChiKnowPo",
        "citation_url": "https://huggingface.co/datasets/calfa-ai/chiknowpo",
    },
]

# Sandbox-ICL stems → which language tag they belong to.
SANDBOX_ICL_SPECS = [
    {"stem": "FRAN_0393_01209_L",           "language": "Middle French — E-NDP"},
    {"stem": "FRAN_0393_01602_L",           "language": "Middle French — E-NDP"},
    {"stem": "BULAC_MS_ARA_6_49775",        "language": "Arabic — Rasam corpus"},
    {"stem": "BULAC_MS_ARA_6_49810",        "language": "Arabic — Rasam corpus"},
    {"stem": "BULAC_BIULO_CHI_1087_1_0435", "language": "Classical Chinese — ChiKnowPo"},
]


def _load_sandbox_into_session() -> None:
    """Load 4 sandbox pages + seed the global pool with 5 sandbox ICL examples.
    Sandbox ICL is tagged with the same language as its target page, so the
    pool filter naturally yields 0/1/2-shot per page. All entries live in
    one shared pool the user can curate (delete, export, add to)."""
    samples_dir = DATA_DIR / "samples"
    icl_dir = DATA_DIR / "ICL"

    if samples_dir.is_dir():
        for spec in SANDBOX_SPECS:
            f = samples_dir / spec["file"]
            if not f.is_file():
                continue
            try:
                raw = f.read_bytes()
                b64, w, h = load_image_bytes_to_jpeg_b64(raw)
            except Exception:
                continue
            image_id = uuid.uuid4().hex[:12]
            SESSION["images"][image_id] = {
                "b64": b64, "filename": spec["file"], "w": w, "h": h,
            }
            page = _new_page(image_id)
            page["title"] = spec["title"]
            page["short_label"] = spec["short"]
            page["language"] = spec["language"]
            page["citation_url"] = spec["citation_url"]
            page["is_sandbox"] = True
            SESSION["pages"].append(page)

    if icl_dir.is_dir() and len(SESSION["icl_pool"]) == 0:
        for spec in SANDBOX_ICL_SPECS:
            img_path = icl_dir / f"{spec['stem']}.jpg"
            txt_path = icl_dir / f"{spec['stem']}.txt"
            if not (img_path.is_file() and txt_path.is_file()):
                continue
            try:
                raw = img_path.read_bytes()
                b64, _, _ = load_image_bytes_to_jpeg_b64(raw)
                text = txt_path.read_text(encoding="utf-8").rstrip()
            except Exception:
                continue
            SESSION["icl_pool"].add(
                image_b64=b64, text=text,
                language=spec["language"], source="sandbox",
            )


def _reset_session_state() -> None:
    """Full wipe: pages, images, ICL pool, totals, mode, AND settings (prompts,
    JSON template, guidelines, language, custom overrides) — back to defaults.
    Sandbox pages and seed ICL are reloaded. The OpenRouter API key lives
    client-side in localStorage and is therefore untouched."""
    SESSION["settings"] = _default_settings()
    SESSION["images"] = {}
    SESSION["pages"] = []
    SESSION["icl_pool"] = ICLPool()
    SESSION["totals"] = {"requests": 0, "input_tokens": 0, "output_tokens": 0, "cost_usd": 0.0}
    SESSION["mode"] = "ocr"
    _load_sandbox_into_session()


# Bootstrap call is deferred to the end of the module, after _new_page() is defined.


def _accumulate(usage: Optional[dict]) -> None:
    """Bump session totals from one OpenRouter usage dict (no-op if usage is None)."""
    if not usage:
        return
    t = SESSION["totals"]
    t["requests"] += 1
    pt = usage.get("prompt_tokens") or 0
    ct = usage.get("completion_tokens") or 0
    cost = usage.get("cost")
    try:
        t["input_tokens"] += int(pt)
        t["output_tokens"] += int(ct)
        if cost is not None:
            t["cost_usd"] = round(float(t["cost_usd"]) + float(cost), 6)
    except (TypeError, ValueError):
        pass


def _sum_usages(*usages: Optional[dict]) -> dict:
    """Aggregate several usage dicts into one (None entries skipped)."""
    out = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "cost": 0.0, "has_cost": False}
    for u in usages:
        if not u: continue
        out["prompt_tokens"] += int(u.get("prompt_tokens") or 0)
        out["completion_tokens"] += int(u.get("completion_tokens") or 0)
        out["total_tokens"] += int(u.get("total_tokens") or 0)
        c = u.get("cost")
        if c is not None:
            out["cost"] = round(float(out["cost"]) + float(c), 6)
            out["has_cost"] = True
    if not out["has_cost"]:
        out["cost"] = None
    out.pop("has_cost")
    return out


def _new_page(image_id: str) -> dict:
    return {
        "id": uuid.uuid4().hex[:10],
        "image_id": image_id,
        "ocr_text": "",
        "corrected_text": "",
        "diff": [],
        "cer": None,
        "wer": None,
        "status": "pending",  # pending | ocr_running | ocr_done | corrected | moe_running | moe_done | error
        "error": "",
        "moe": None,           # last MoE result dict (post-correction mode)
        "title": "",           # full display label (sandbox pages set this)
        "short_label": "",     # short sidebar label (sandbox pages set this)
        "language": "",        # per-page override for settings["language"]
        "citation_url": "",    # external dataset URL (sandbox pages)
        "is_sandbox": False,
        # Object-detection results (zero-shot via grounding-capable VLM).
        # Each item: {"id": str, "label": str, "bbox_px": [x1,y1,x2,y2],
        #             "box_2d": [ymin,xmin,ymax,xmax] (0-1000), "confidence": float}
        "detections": [],
        "detect_query": "",
        "detect_model": "",
        "detect_raw": "",
    }


def _public_image(img: dict) -> dict:
    return {"filename": img["filename"], "w": img.get("w", 0), "h": img.get("h", 0)}


def _public_state() -> dict:
    return {
        "settings": SESSION["settings"],
        "images": {k: _public_image(v) for k, v in SESSION["images"].items()},
        "pages": [dict(p) for p in SESSION["pages"]],
        "icl_pool": SESSION["icl_pool"].public_view(),
        "mode": SESSION["mode"],
        "curated_models": CURATED_MODELS,
        "env_key_present": bool(ENV_OPENROUTER_KEY),
        "totals": SESSION["totals"],
    }


def _resolve_key(header_key: Optional[str]) -> str:
    return (header_key or "").strip() or ENV_OPENROUTER_KEY


app = FastAPI(title="HTR VLM + MoE")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
    allow_credentials=False,
)
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


@app.get("/")
def index():
    return FileResponse(str(STATIC_DIR / "index.html"))


@app.get("/api/state")
def get_state():
    return _public_state()


@app.get("/api/image/{image_id}")
def get_image(image_id: str):
    img = SESSION["images"].get(image_id)
    if not img:
        raise HTTPException(404, "Image not found")
    return {"image_id": image_id, "b64": img["b64"], "filename": img["filename"]}


@app.post("/api/settings")
def post_settings(req: SettingsReq):
    SESSION["settings"].update(req.model_dump())
    return _public_state()


@app.post("/api/mode/{mode}")
def post_mode(mode: str):
    if mode not in ("ocr", "post_correction"):
        raise HTTPException(400, "Invalid mode")
    SESSION["mode"] = mode
    return _public_state()


@app.post("/api/settings/test_key")
def post_test_key(req: TestKeyReq, x_api_key: Optional[str] = Header(default=None)):
    key = _resolve_key(x_api_key)
    if not key:
        return {"ok": False, "message": "No API key provided."}
    ok, msg = test_connection_sync(key, model=req.model)
    return {"ok": ok, "message": msg}


MAX_UPLOAD_BYTES = 10 * 1024 * 1024  # 10 MB per image — abuse cap for public deploys.


@app.post("/api/upload")
async def post_upload(files: list[UploadFile] = File(...)):
    out_pages: list[str] = []
    for f in files:
        raw = await f.read()
        if not raw:
            continue
        if len(raw) > MAX_UPLOAD_BYTES:
            raise HTTPException(
                413,
                f"{f.filename!r} is {len(raw) / 1_048_576:.1f} MB; per-image limit is "
                f"{MAX_UPLOAD_BYTES // 1_048_576} MB.",
            )
        try:
            b64, w, h = load_image_bytes_to_jpeg_b64(raw)
        except Exception as e:  # noqa: BLE001
            raise HTTPException(400, f"Could not load image {f.filename!r}: {e}")
        image_id = uuid.uuid4().hex[:12]
        SESSION["images"][image_id] = {"b64": b64, "filename": f.filename or image_id, "w": w, "h": h}
        page = _new_page(image_id)
        SESSION["pages"].append(page)
        out_pages.append(page["id"])
    return {"created_pages": out_pages, "state": _public_state()}


@app.delete("/api/page/{page_idx}")
def delete_page(page_idx: int):
    if page_idx < 0 or page_idx >= len(SESSION["pages"]):
        raise HTTPException(404, "Page out of range")
    page = SESSION["pages"].pop(page_idx)
    img_id = page["image_id"]
    if img_id and not any(p["image_id"] == img_id for p in SESSION["pages"]):
        SESSION["images"].pop(img_id, None)
    return _public_state()


@app.post("/api/ocr")
async def post_ocr(req: OcrReq, x_api_key: Optional[str] = Header(default=None)):
    key = _resolve_key(x_api_key)
    if not key:
        raise HTTPException(400, "Missing OpenRouter API key (X-API-Key header or OPENROUTER_API_KEY env).")
    if req.page_idx < 0 or req.page_idx >= len(SESSION["pages"]):
        raise HTTPException(404, "Page out of range")

    page = SESSION["pages"][req.page_idx]
    image = SESSION["images"].get(page["image_id"])
    if not image:
        raise HTTPException(404, "Image missing for that page")

    settings = SESSION["settings"]
    page["status"] = "ocr_running"
    page["error"] = ""

    # Per-page language wins over the global setting (sandbox pages ship with
    # their own tag so they pick the right ICL slice out of the box).
    language = (page.get("language") or "").strip() or settings["language"]

    output_mode = settings.get("output_mode") or DEFAULT_OUTPUT_MODE
    system = render_ocr_system(
        language=language,
        guidelines=settings["guidelines"],
        override=settings.get("custom_ocr_system"),
        mode=output_mode,
        json_template=settings.get("output_json_template"),
    )

    icl_examples: list = []
    if settings.get("use_icl", True):
        n = max(0, int(settings.get("n_icl", 0) or 0))
        if n:
            icl_examples = SESSION["icl_pool"].sample(n, language=language)

    user_parts: list = []
    if icl_examples:
        user_parts.append(text_part(render_ocr_few_shot_header(len(icl_examples))))
        for i, ex in enumerate(icl_examples, 1):
            user_parts.append(text_part(f"\n--- Example {i} (image) ---"))
            user_parts.append(image_part(ex.image_b64))
            user_parts.append(text_part(f"--- Example {i} (validated transcription) ---\n{ex.text}\n"))
    user_parts.append(text_part(render_ocr_user_target(few_shot=bool(icl_examples), override=settings.get("custom_ocr_user"))))
    user_parts.append(image_part(image["b64"]))

    async with LLMClient(api_key=key) as client:
        res = await client.chat(
            model=settings["ocr_model"],
            system=system,
            user_parts=user_parts,
            temperature=float(settings.get("temperature", 0.0)),
            max_tokens=4096,
            force_json=True,
        )

    if not res.ok:
        page["status"] = "error"
        page["error"] = res.error
        raise HTTPException(502, f"OCR call failed: {res.error}")

    _accumulate(res.usage)
    lines, structured = parse_lines_from_model_response(res.content, mode=output_mode)
    text = "\n".join(lines).rstrip()
    page["ocr_text"] = text
    page["ocr_structured"] = structured  # may be None for plain "lines" mode
    page["ocr_mode"] = output_mode
    page["corrected_text"] = ""
    page["diff"] = []
    page["cer"] = None
    page["wer"] = None
    page["status"] = "ocr_done"
    page["usage"] = res.usage
    return {
        "page_idx": req.page_idx,
        "ocr_text": text,
        "ocr_structured": structured,
        "ocr_mode": output_mode,
        "raw": res.content[:1000],
        "latency_s": res.latency_s,
        "usage": res.usage,
        "state": _public_state(),
    }


@app.post("/api/correction")
def post_correction(req: CorrectionReq):
    if req.page_idx < 0 or req.page_idx >= len(SESSION["pages"]):
        raise HTTPException(404, "Page out of range")
    page = SESSION["pages"][req.page_idx]
    ref = (page.get("ocr_text") or "").rstrip()
    corrected = (req.corrected_text or "").rstrip()
    page["corrected_text"] = corrected
    page["diff"] = diff_words(ref, corrected)
    metrics = compute_cer_wer(reference=corrected, hypothesis=ref)
    page["cer"] = metrics["cer"]
    page["wer"] = metrics["wer"]
    page["status"] = "corrected"
    added = None
    if req.add_to_icl and corrected:
        image = SESSION["images"].get(page["image_id"])
        if image:
            ex = SESSION["icl_pool"].add(
                image_b64=image["b64"],
                text=corrected,
                language=SESSION["settings"]["language"],
                source="corrected",
            )
            added = ex.id
    return {
        "page_idx": req.page_idx,
        "diff": page["diff"],
        "cer": page["cer"],
        "wer": page["wer"],
        "icl_added": added,
        "state": _public_state(),
    }


@app.post("/api/icl/paste")
def post_icl_paste(req: PasteIclReq):
    image = SESSION["images"].get(req.image_id)
    if not image:
        raise HTTPException(404, "Image not found")
    ex = SESSION["icl_pool"].add(
        image_b64=image["b64"],
        text=req.text.rstrip(),
        language=req.language or SESSION["settings"]["language"],
        source="pasted",
    )
    return {"id": ex.id, "state": _public_state()}


@app.post("/api/icl/remove/{item_id}")
def post_icl_remove(item_id: str):
    ok = SESSION["icl_pool"].remove(item_id)
    if not ok:
        raise HTTPException(404, "ICL item not found")
    return _public_state()


@app.get("/api/icl/export")
def get_icl_export():
    items = SESSION["icl_pool"].to_jsonl_dicts()
    text = export_icl_jsonl(items)
    return PlainTextResponse(text, media_type="application/x-ndjson")


@app.post("/api/post_correct")
async def post_post_correct(req: PostCorrReq, x_api_key: Optional[str] = Header(default=None)):
    key = _resolve_key(x_api_key)
    if not key:
        raise HTTPException(400, "Missing OpenRouter API key")
    if not (req.ocr_text or "").strip():
        raise HTTPException(400, "ocr_text is empty")

    image_b64 = None
    image_id = req.image_id
    page_idx = req.page_idx
    if page_idx is not None and 0 <= page_idx < len(SESSION["pages"]):
        page = SESSION["pages"][page_idx]
        image_id = image_id or page["image_id"]
    if image_id and image_id in SESSION["images"]:
        image_b64 = SESSION["images"][image_id]["b64"]
    if not image_b64:
        raise HTTPException(400, "An image is required for post-correction (upload one or run OCR first).")

    settings = SESSION["settings"]
    if page_idx is not None and 0 <= page_idx < len(SESSION["pages"]):
        SESSION["pages"][page_idx]["status"] = "moe_running"

    result = await run_moe_correction(
        api_key=key,
        image_b64=image_b64,
        ocr_text=req.ocr_text,
        expert_a_model=settings["moe_expert_a"],
        expert_b_model=settings["moe_expert_b"],
        judge_model=settings["moe_judge"],
        language=settings["language"],
        guidelines=settings["guidelines"],
        temperature=float(settings.get("temperature", 0.1)),
        expert_system_override=settings.get("custom_expert_system") or None,
        expert_user_override=settings.get("custom_expert_user") or None,
        judge_system_override=settings.get("custom_judge_system") or None,
        judge_user_override=settings.get("custom_judge_user") or None,
    )

    diff = diff_words(req.ocr_text, result["final_text"])
    metrics = compute_cer_wer(reference=result["final_text"], hypothesis=req.ocr_text)

    ea_usage = result["expert_a"].get("usage")
    eb_usage = result["expert_b"].get("usage")
    judge_usage = (result["judge"] or {}).get("usage")
    for u in (ea_usage, eb_usage, judge_usage):
        _accumulate(u)
    total_usage = _sum_usages(ea_usage, eb_usage, judge_usage)

    payload = {
        "final_text": result["final_text"],
        "source": result["source"],
        "confidence": result["confidence"],
        "rationale": result["rationale"],
        "diff": diff,
        "cer": metrics["cer"],
        "wer": metrics["wer"],
        "usage_total": total_usage,
        "expert_a": {k: result["expert_a"][k] for k in ("model", "ok", "error", "latency_s", "corrected_text", "confidence", "corrections", "usage")},
        "expert_b": {k: result["expert_b"][k] for k in ("model", "ok", "error", "latency_s", "corrected_text", "confidence", "corrections", "usage")},
        "judge": ({k: result["judge"][k] for k in ("model", "ok", "error", "latency_s", "final_text", "confidence", "source", "rationale", "usage")} if result["judge"] else None),
    }

    if page_idx is not None and 0 <= page_idx < len(SESSION["pages"]):
        page = SESSION["pages"][page_idx]
        page["moe"] = payload                  # payload does NOT contain `state` yet
        page["corrected_text"] = result["final_text"]
        page["diff"] = diff
        page["cer"] = metrics["cer"]
        page["wer"] = metrics["wer"]
        page["status"] = "moe_done"

    # Build the response as a fresh dict so we never mutate the version stored
    # in page["moe"] (mutating would create a state → pages[i].moe → state cycle
    # that crashes jsonable_encoder).
    return {**payload, "state": _public_state()}


@app.get("/api/page/{page_idx}/moe")
def get_page_moe(page_idx: int):
    if page_idx < 0 or page_idx >= len(SESSION["pages"]):
        raise HTTPException(404, "Page out of range")
    page = SESSION["pages"][page_idx]
    if page.get("moe") is None:
        raise HTTPException(404, "No MoE result for this page")
    return page["moe"]


# ───────────────────────────────────────────────────────────────────────────
#  Click-to-annotate via FastSAM (local, point prompt → bbox)
# ───────────────────────────────────────────────────────────────────────────

# Stable colour palette so the same label always renders in the same colour.
_DETECT_PALETTE = [
    "#e11d48",  # rose-600
    "#0ea5e9",  # sky-500
    "#16a34a",  # green-600
    "#a855f7",  # purple-500
    "#ea580c",  # orange-600
    "#0891b2",  # cyan-600
    "#65a30d",  # lime-600
    "#db2777",  # pink-600
]


def _color_for(label: str) -> str:
    h = sum(ord(c) for c in (label or "?")) % len(_DETECT_PALETTE)
    return _DETECT_PALETTE[h]


def _next_annotation_id(page: dict) -> str:
    used = {d.get("id") for d in (page.get("detections") or [])}
    i = 1
    while f"ann-{i:03d}" in used:
        i += 1
    return f"ann-{i:03d}"


@app.post("/api/sam/point")
def post_sam_point(req: SamPointReq):
    """Click-to-annotate: feed (x, y) to FastSAM, return the bbox of the
    smallest mask containing that pixel. Stores the new annotation on the page."""
    if req.page_idx < 0 or req.page_idx >= len(SESSION["pages"]):
        raise HTTPException(404, "Page out of range")
    page = SESSION["pages"][req.page_idx]
    image = SESSION["images"].get(page["image_id"])
    if not image:
        raise HTTPException(404, "Image missing for that page")

    try:
        from detector import segment_at_point
    except ImportError as e:
        raise HTTPException(500, f"FastSAM unavailable: {e}. Run `pip install ultralytics`.")

    result = segment_at_point(image["b64"], int(req.x), int(req.y))
    if not result:
        raise HTTPException(404, "No object found at that point. Try clicking closer to the object centre.")

    label = (req.label or "").strip() or "object"
    ann = {
        "id": _next_annotation_id(page),
        "label": label,
        "bbox_px": result["bbox_px"],
        "confidence": None,         # FastSAM masks have no confidence per-instance
        "color": _color_for(label),
        "source": "fastsam",
        "click_px": [int(req.x), int(req.y)],
    }
    page.setdefault("detections", []).append(ann)
    return {"page_idx": req.page_idx, "annotation": ann, "state": _public_state()}


@app.post("/api/annotation/label")
def post_relabel(req: AnnotationLabelReq):
    if req.page_idx < 0 or req.page_idx >= len(SESSION["pages"]):
        raise HTTPException(404, "Page out of range")
    page = SESSION["pages"][req.page_idx]
    new_label = (req.label or "").strip() or "object"
    for d in (page.get("detections") or []):
        if d.get("id") == req.annotation_id:
            d["label"] = new_label
            d["color"] = _color_for(new_label)
            return {"page_idx": req.page_idx, "annotation": d, "state": _public_state()}
    raise HTTPException(404, "Annotation not found")


@app.delete("/api/page/{page_idx}/annotation/{annotation_id}")
def delete_annotation(page_idx: int, annotation_id: str):
    if page_idx < 0 or page_idx >= len(SESSION["pages"]):
        raise HTTPException(404, "Page out of range")
    page = SESSION["pages"][page_idx]
    before = len(page.get("detections") or [])
    page["detections"] = [d for d in (page.get("detections") or []) if d.get("id") != annotation_id]
    if len(page["detections"]) == before:
        raise HTTPException(404, "Annotation not found")
    return _public_state()


@app.get("/api/ocr/preview/{page_idx}")
def preview_ocr_prompt(page_idx: int):
    """Dry-run: render the exact system + user message parts that would be sent
    to OpenRouter for OCR on this page, WITHOUT actually calling the model.
    Image bytes are stripped and replaced by descriptive placeholders so the
    response stays small.
    """
    if page_idx < 0 or page_idx >= len(SESSION["pages"]):
        raise HTTPException(404, "Page out of range")
    page = SESSION["pages"][page_idx]
    image = SESSION["images"].get(page["image_id"])
    if not image:
        raise HTTPException(404, "Image missing for that page")

    settings = SESSION["settings"]
    language = (page.get("language") or "").strip() or settings["language"]
    output_mode = settings.get("output_mode") or DEFAULT_OUTPUT_MODE
    system = render_ocr_system(
        language=language,
        guidelines=settings["guidelines"],
        override=settings.get("custom_ocr_system"),
        mode=output_mode,
        json_template=settings.get("output_json_template"),
    )

    icl_examples: list = []
    if settings.get("use_icl", True):
        n = max(0, int(settings.get("n_icl", 0) or 0))
        if n:
            icl_examples = SESSION["icl_pool"].sample(n, language=language)

    # Build the parts WITHOUT base64 — return descriptive placeholders only.
    parts: list[dict] = []
    if icl_examples:
        parts.append({"type": "text", "text": render_ocr_few_shot_header(len(icl_examples))})
        for i, ex in enumerate(icl_examples, 1):
            parts.append({"type": "text", "text": f"\n--- Example {i} (image) ---"})
            parts.append({
                "type": "image",
                "label": f"<ICL example {i} image · id={ex.id} · {len(ex.image_b64)} b64 chars>",
            })
            parts.append({"type": "text", "text": f"--- Example {i} (validated transcription) ---\n{ex.text}\n"})
    parts.append({
        "type": "text",
        "text": render_ocr_user_target(
            few_shot=bool(icl_examples),
            override=settings.get("custom_ocr_user"),
        ),
    })
    parts.append({
        "type": "image",
        "label": f"<TARGET PAGE image · {image.get('filename', '')} · {image.get('w','?')}×{image.get('h','?')}>",
    })

    return {
        "page_idx": page_idx,
        "model": settings["ocr_model"],
        "output_mode": output_mode,
        "icl": {
            "enabled": bool(settings.get("use_icl", True)),
            "n_requested": int(settings.get("n_icl", 0) or 0),
            "n_injected": len(icl_examples),
            "language_filter": language or None,
            "pool_size": len(SESSION["icl_pool"]),
        },
        "system": system,
        "user_parts": parts,
    }


# ───────────────────────────────────────────────────────────────────────────
#  Prompts defaults / presets
# ───────────────────────────────────────────────────────────────────────────

@app.get("/api/prompts/defaults")
def get_prompts_defaults():
    """Expose default prompt templates so the UI can show 'using default'."""
    return {
        "defaults": {
            "ocr_system": OCR_SYSTEM_TPL,
            "ocr_user": OCR_USER_ZERO_SHOT_TPL,
            "expert_system": EXPERT_SYSTEM_TPL,
            "expert_user": EXPERT_USER_TPL,
            "judge_system": JUDGE_SYSTEM_TPL,
            "judge_user": JUDGE_USER_TPL,
        },
        "preset_names": list_preset_names(),
    }


@app.get("/api/prompts/preset")
def get_prompt_preset(family: str, name: str):
    body = get_preset(family, name)
    if not body:
        raise HTTPException(404, f"Preset {family}/{name} not found")
    return {"family": family, "name": name, "body": body}


# ───────────────────────────────────────────────────────────────────────────
#  Exports
# ───────────────────────────────────────────────────────────────────────────

def _safe_filename(s: str) -> str:
    return "".join(c if c.isalnum() or c in ("-", "_", ".") else "_" for c in s)[:80]


@app.get("/api/page/{page_idx}/export")
def export_page(page_idx: int, fmt: str = "txt"):
    if page_idx < 0 or page_idx >= len(SESSION["pages"]):
        raise HTTPException(404, "Page out of range")
    page = SESSION["pages"][page_idx]
    image = SESSION["images"].get(page["image_id"]) or {}
    stem = _safe_filename((image.get("filename") or f"page_{page_idx+1}").rsplit(".", 1)[0])

    if fmt == "txt":
        text = page.get("corrected_text") or page.get("ocr_text") or ""
        return PlainTextResponse(
            text,
            media_type="text/plain; charset=utf-8",
            headers={"Content-Disposition": f'attachment; filename="{stem}.txt"'},
        )
    if fmt == "json":
        payload = {
            "filename": image.get("filename"),
            "image_w": image.get("w"),
            "image_h": image.get("h"),
            "ocr_text": page.get("ocr_text"),
            "corrected_text": page.get("corrected_text"),
            "cer": page.get("cer"),
            "wer": page.get("wer"),
            "status": page.get("status"),
            "moe": page.get("moe"),  # full MoE detail if any
            "detections": page.get("detections") or [],
            "detect_query": page.get("detect_query") or "",
            "detect_model": page.get("detect_model") or "",
        }
        return JSONResponse(
            payload,
            headers={"Content-Disposition": f'attachment; filename="{stem}.json"'},
        )
    raise HTTPException(400, f"Unsupported format {fmt!r}; use txt or json")


@app.get("/api/pages/export")
def export_pages_all(fmt: str = "jsonl"):
    rows = []
    for idx, p in enumerate(SESSION["pages"]):
        img = SESSION["images"].get(p["image_id"]) or {}
        rows.append({
            "idx": idx,
            "filename": img.get("filename"),
            "image_w": img.get("w"),
            "image_h": img.get("h"),
            "image_b64": img.get("b64"),
            "ocr_text": p.get("ocr_text") or "",
            "corrected_text": p.get("corrected_text") or "",
            "cer": p.get("cer"),
            "wer": p.get("wer"),
            "moe_final_text": (p.get("moe") or {}).get("final_text"),
            "moe_source": (p.get("moe") or {}).get("source"),
            "detections": p.get("detections") or [],
            "detect_query": p.get("detect_query") or "",
        })
    if fmt == "jsonl":
        text = "\n".join(json.dumps(r, ensure_ascii=False) for r in rows)
        return PlainTextResponse(
            text,
            media_type="application/x-ndjson",
            headers={"Content-Disposition": 'attachment; filename="pages_export.jsonl"'},
        )
    if fmt == "txt":
        text = "\n\n".join(
            f"=== Page {r['idx']+1} · {r['filename']} ===\n{r['corrected_text'] or r['ocr_text']}"
            for r in rows
        )
        return PlainTextResponse(
            text,
            media_type="text/plain; charset=utf-8",
            headers={"Content-Disposition": 'attachment; filename="pages_export.txt"'},
        )
    raise HTTPException(400, f"Unsupported format {fmt!r}; use jsonl or txt")


# ───────────────────────────────────────────────────────────────────────────
#  Sandbox: pre-loaded sample images for first-time users
# ───────────────────────────────────────────────────────────────────────────

@app.get("/api/sandbox/list")
def list_sandbox_samples():
    samples_dir = DATA_DIR / "samples"
    if not samples_dir.is_dir():
        return {"samples": []}
    files = sorted([
        p.name for p in samples_dir.iterdir()
        if p.is_file() and p.suffix.lower() in (".jpg", ".jpeg", ".png", ".webp")
    ])
    return {"samples": files}


@app.post("/api/session/reset")
def reset_session():
    """Full wipe: pages, images, ICL pool, cost totals, AND settings
    (prompts, JSON template, guidelines, language, custom overrides) — back
    to defaults. Sandbox pages and seed ICL are re-loaded."""
    _reset_session_state()
    return _public_state()


# ───────────────────────────────────────────────────────────────────────────
#  Citation
# ───────────────────────────────────────────────────────────────────────────

CITATION_BIB = """@software{vidal-gorene_htr_vlm_2026,
  author  = {Vidal-Gorène, Chahan},
  title   = {{HTR · VLM Annotator}: a custom HuggingFace app for HTR by VLMs with MoE post-correction},
  year    = {2026},
  url     = {https://huggingface.co/spaces/CVidalG/htr-vlm-annotator},
  note    = {LLM-as-annotator workflow with two VLM experts + judge}
}"""


@app.get("/api/citation")
def get_citation():
    return {
        "title": "HTR · VLM Annotator",
        "author": "Vidal-Gorène, Chahan",
        "year": 2026,
        "bibtex": CITATION_BIB,
        "apa": (
            "Vidal-Gorène, C. (2026). HTR · VLM Annotator: a custom HuggingFace app "
            "for HTR by VLMs with MoE post-correction. [Software]."
        ),
    }


# Bootstrap: load sandbox pages + the placeholder ICL example on startup.
# Must run after _new_page() (defined above) and after FastAPI routes are
# registered so that an immediate GET /api/state already returns the samples.
_load_sandbox_into_session()


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)