"""FastAPI app for CPU-only, fine-tuned dense retrieval."""

from __future__ import annotations

import csv
import io
import json
import os
import threading
import uuid
from datetime import datetime, timezone
from functools import lru_cache
from pathlib import Path

from fastapi import Body, FastAPI, HTTPException, Response
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles

_NO_CACHE = "no-cache, must-revalidate"


class NoCacheStaticFiles(StaticFiles):
    """Serve the vanilla frontend with revalidation so a rebuilt app.js/styles.css
    is never served stale from the browser cache (bit us during development)."""

    async def get_response(self, path, scope):
        resp = await super().get_response(path, scope)
        resp.headers["Cache-Control"] = _NO_CACHE
        return resp

from .codesearch import CodeSearchEngine, make_embedder
from .encode import EncodeEngine
from .graph import KnowledgeGraph
from . import models as model_registry
from . import planner as query_planner
from .paths import ANNOT_DIR as ANNOT
from .paths import CORPUS_STATS, COVERAGE_REPORT, LAB_NOISE_VOCAB
from .retriever import DenseEmbedder

FRONTEND = Path(__file__).resolve().parents[1] / "frontend"

app = FastAPI(title="ENCODE: Clinical Code and Phenotype Search")


@lru_cache(maxsize=1)
def _embedder() -> DenseEmbedder:
    return make_embedder()


@lru_cache(maxsize=1)
def codes() -> CodeSearchEngine:
    return CodeSearchEngine(_embedder())


# One engine per prebuilt vector set, keyed on the directory so the default
# model is never loaded twice. Selecting another phenotype model is a registry
# entry pointing at its own embeddings; each set records its own pooling and
# prefix convention in its config.json.
#
# Construction parses a ~190MB phenotype file and loads its own copy of the
# query model. On the Space that data sits on a network bucket, so a build
# takes minutes; it must never run on a visitor's request (the Netlify proxy
# times out long before it finishes). The startup warm thread builds the
# default engine; a request that arrives first gets a clear 503 instead.
_engines: dict[str, EncodeEngine] = {}
_ENGINE_LOCK = threading.Lock()


def engine(emb_dir: str) -> EncodeEngine:
    made = _engines.get(emb_dir)
    if made is not None:
        return made
    if not _ENGINE_LOCK.acquire(blocking=False):
        raise HTTPException(503, "The phenotype index is still loading. Try again in a minute.")
    try:
        if emb_dir not in _engines:
            _engines[emb_dir] = EncodeEngine(emb_dir=emb_dir)
        return _engines[emb_dir]
    finally:
        _ENGINE_LOCK.release()


def _warm_phenotype() -> None:
    try:
        with _ENGINE_LOCK:
            emb_dir = str(model_registry.resolve(None, "phenotype").pheno_emb_dir)
            if emb_dir not in _engines:
                _engines[emb_dir] = EncodeEngine(emb_dir=emb_dir)
    except Exception as err:  # warm failure must not kill the server
        print(f"phenotype warm failed: {err}", flush=True)


def _pheno_engine(spec: model_registry.ModelSpec) -> EncodeEngine:
    return engine(str(spec.pheno_emb_dir))


def _default_pheno_engine() -> EncodeEngine:
    """Detail lookups (phenotype record, code hierarchy) are model-independent —
    they read runtime phenotype metadata, not vectors — so they use the default build."""
    return _pheno_engine(_resolve_model(None, "phenotype"))


def _resolve_model(model_id: str | None, category: str) -> model_registry.ModelSpec:
    try:
        return model_registry.resolve(model_id, category)
    except model_registry.ModelError as err:
        raise HTTPException(err.status, err.detail)


def _stamp(payload: dict, spec: model_registry.ModelSpec) -> dict:
    """Every result set says which model produced it."""
    payload["model_id"] = spec.id
    payload["model_label"] = spec.label
    # The engine's own payload names the model by its filesystem path; the
    # response should carry the label, matching what code search reports.
    payload["model"] = spec.label
    return payload


# One result cap for every retrieval endpoint, mirrored by the count box in
# the sidebar. Lab reviews legitimately run to thousands of rows.
K_MAX = 2000


def _k(k: int) -> int:
    return min(max(k, 1), K_MAX)


@lru_cache(maxsize=1)
def graph() -> KnowledgeGraph:
    return KnowledgeGraph(diagnosis_records=codes().records("diagnosis"),
                          procedure_records=codes().records("procedure"))


@app.on_event("startup")
def _warm() -> None:
    codes()   # load the fine-tuned model; FAISS indexes stay lazy by category
    threading.Thread(target=_warm_phenotype, daemon=True).start()


@app.get("/healthz")
def health() -> dict:
    return {"status": "ok", "device": _embedder().device}


# Uptime monitors commonly probe with HEAD, which this app otherwise answers
# with 404: FastAPI does not map HEAD onto GET routes here, so the two probe
# targets get explicit handlers.
@app.head("/healthz")
def health_head() -> Response:
    return Response(status_code=200)


@app.head("/")
def root_head() -> Response:
    return Response(status_code=200)


# -- code search (primary) -------------------------------------------------

@app.get("/api/code/categories")
def code_categories() -> dict:
    return {"categories": codes().categories()}


@app.get("/api/models")
def model_catalog() -> dict:
    """Retrieval models this deployment knows about, and which ones it serves."""
    return model_registry.catalog()


@app.get("/api/corpus")
def corpus() -> dict:
    """Index sizes and mapping coverage for the data release now loaded.

    The numbers come from scripts/report_coverage.py, which is the only
    generator of coverage statistics in this project; this endpoint serves
    what that report wrote, so the About panel cannot state a figure the
    standing report does not. A deployment without the file simply has no
    corpus section."""
    for path in (CORPUS_STATS, COVERAGE_REPORT):
        if path.exists():
            return json.loads(path.read_text(encoding="utf-8"))
    raise HTTPException(404, "No coverage report in this deployment")


@app.get("/api/lab/noise")
def lab_noise_vocab() -> dict:
    """The mined lab merge-noise vocabulary, with its evidence.

    Written by scripts/build_lab_noise_vocab.py, the only generator, so the
    merge panel cannot show a word the miner did not learn. A deployment
    without the file merges on the fixed rules only."""
    if LAB_NOISE_VOCAB.exists():
        return json.loads(LAB_NOISE_VOCAB.read_text(encoding="utf-8"))
    raise HTTPException(404, "No noise vocabulary in this deployment")


# -- query planner -----------------------------------------------------------
# Decomposes one natural-language cohort description into search criteria. The
# only path in this application that sends user text off the deployment: the
# query string goes to DeepSeek, nothing else. No search results, no
# annotations, no collected codes, and no conversation history are included.


@app.get("/api/plan/status")
def plan_status() -> dict:
    """What planner model, if any, this deployment ships. A user who adds their
    own model can plan even when `available` is false, so the frontend decides
    whether to offer the mode from this plus its own saved models."""
    catalog = query_planner.builtin_catalog()
    return {"available": bool(catalog),
            # `models` is the picker's list, default first. `model` is the
            # default's label, kept for a frontend that predates the list.
            "models": [{"id": m["id"], "label": m["label"]} for m in catalog],
            "model": catalog[0]["label"] if catalog else None,
            "formats": list(query_planner.KINDS),
            # The browser needs the prompt to call its own model directly.
            # Serving it keeps one copy of the instructions, in planner.py.
            "prompt": query_planner.SYSTEM}


@app.post("/api/plan")
def plan(payload: dict = Body(...)) -> dict:
    """`model` optionally carries a user-supplied provider
    ({kind, base_url, model, api_key, label}). Those credentials belong to the
    caller: they are used for one outbound call and never stored or logged."""
    q = (payload.get("q") or "").strip()
    if not q:
        raise HTTPException(400, "Empty query")
    try:
        return query_planner.plan(q, payload.get("model"), payload.get("builtin"))
    except query_planner.LlmError as err:
        # 503, not 500: this is an upstream/config outage, and the UI tells the
        # user to use the regular search rather than implying a bad query.
        raise HTTPException(503, str(err))


@app.post("/api/plan/stream")
def plan_stream(payload: dict = Body(...)):
    """Server-sent events for the built-in model: the reasoning as it happens,
    then the validated plan.

    POST rather than GET/EventSource because the description can be long, and
    a URL is the wrong place for a clinical query. The frontend reads the body
    as a stream and parses the SSE frames itself."""
    q = (payload.get("q") or "").strip()
    if not q:
        raise HTTPException(400, "Empty query")

    def frames():
        try:
            for kind, value in query_planner.plan_streaming(q, payload.get("builtin")):
                if kind == "thinking":
                    yield f"data: {json.dumps({'type': 'thinking', 'text': value})}\n\n"
                elif kind == "usage":
                    yield f"data: {json.dumps({'type': 'usage', 'usage': value})}\n\n"
                else:
                    yield f"data: {json.dumps({'type': 'plan', 'plan': value})}\n\n"
        except query_planner.LlmError as err:
            # The response has already begun, so an error is a frame, not a
            # status code; the client reports it the same either way.
            yield f"data: {json.dumps({'type': 'error', 'message': str(err)})}\n\n"

    return StreamingResponse(frames(), media_type="text/event-stream",
                             headers={"Cache-Control": _NO_CACHE,
                                      "X-Accel-Buffering": "no"})


@app.post("/api/plan/validate")
def plan_validate(payload: dict = Body(...)) -> dict:
    """Turn a model reply the *browser* obtained into a validated plan.

    This is the path for a user's own model: their browser calls the provider
    directly, so no base URL, model name, or API key is ever sent here. What
    arrives is the query and the model's answer, and every schema rule runs
    server-side exactly as it does for the built-in model."""
    q = (payload.get("q") or "").strip()
    if not q:
        raise HTTPException(400, "Empty query")
    try:
        return query_planner.plan_from_text(q, payload.get("text"), payload.get("label"))
    except query_planner.LlmError as err:
        raise HTTPException(400, str(err))


@app.get("/api/code/systems")
def code_systems(category: str) -> dict:
    try:
        return {"category": category, "systems": codes().systems(category)}
    except KeyError:
        raise HTTPException(404, f"Unknown category '{category}'")


@app.get("/api/code/search")
def code_search(category: str, q: str, k: int = 50, model: str | None = None,
                systems: str | None = None) -> dict:
    if not q.strip():
        raise HTTPException(400, "Empty query")
    spec = _resolve_model(model, category)
    chosen = {s.strip() for s in (systems or "").split(",") if s.strip()} or None
    try:
        return _stamp(codes().search(category, q, k=_k(k),
                                     systems=chosen), spec)
    except KeyError:
        raise HTTPException(404, f"Unknown category '{category}'")


@app.get("/api/code/lookup")
def code_lookup(category: str, code: str, k: int = 50) -> dict:
    """Exact code lookup. No model: nothing here is embedded or ranked."""
    if not code.strip():
        raise HTTPException(400, "Empty code")
    try:
        return codes().lookup(category, code, k=_k(k))
    except KeyError:
        raise HTTPException(404, f"Unknown category '{category}'")


@app.get("/api/code/export")
def code_export(category: str, q: str, k: int = 50, model: str | None = None):
    if not q.strip():
        raise HTTPException(400, "Empty query")
    _resolve_model(model, category)
    try:
        data = codes().search(category, q, k=_k(k))
    except KeyError:
        raise HTTPException(404, f"Unknown category '{category}'")
    buf = io.StringIO()
    w = csv.writer(buf)
    w.writerow(["rank", "code_type", "code", "description", "relevance"])
    for r in data["results"]:
        w.writerow([r["rank"], r["code_type"], r["code"], r["description"], r["relevance"]])
    buf.seek(0)
    fname = f"encode_{category}_{q.strip().replace(' ', '_')[:30]}.csv"
    return StreamingResponse(iter([buf.getvalue()]), media_type="text/csv",
                             headers={"Content-Disposition": f'attachment; filename="{fname}"'})


# -- annotation storage -----------------------------------------------------
# Labels are an append-only record. Nothing here reads, edits, or replaces a
# stored submission: labelling the same query twice, under the same name and
# against the same model, produces two submissions, and both are kept. They
# are told apart by `submitted_at` and by `submission_id`, which is what the
# analysis reads to take the latest labels without losing the earlier ones.
#
# Every submission is written twice under ANNOT:
#
#   <store>.jsonl                          the rolling log the export reads
#   submissions/<kind>/<stamp>_<id>.json   one immutable file per submission
#
# The per-submission file is what makes the record recoverable. An interrupted
# append can leave the rolling log short a line; the individual files still
# hold that submission, and the export reads them back in. They are created
# with mode "x", so no later submission can ever land on top of an earlier
# one. Writes are serialized and flushed to disk, so submissions arriving
# together interleave as whole lines rather than partial ones.

_ANNOT_STORES = {
    "code": ("code_annotations.jsonl",
             ["submission_id", "submitted_at", "annotator", "model", "category", "query"],
             ["rank", "code_type", "code", "description",
              "relevant", "related", "not_relevant", "unsure", "score"]),
    "phenotype": ("query_phenotype_gold.jsonl",
                  ["submission_id", "submitted_at", "annotator", "model", "query"],
                  ["phenotype_id", "title",
                   "relevant", "related", "not_relevant", "unsure", "score"]),
}
_ANNOT_LOCK = threading.Lock()


def _record_submission(kind: str, row: dict) -> dict:
    """Persist one submission and return it, stamped with its own identity."""
    now = datetime.now(timezone.utc)
    # Milliseconds, not seconds: two submissions can land inside the same
    # second, and the timestamp is what orders them.
    stamped = {"submission_id": uuid.uuid4().hex[:12],
               "submitted_at": now.isoformat(timespec="milliseconds"),
               "kind": kind, **row}
    versions = ANNOT / "submissions" / kind
    versions.mkdir(parents=True, exist_ok=True)
    # Sorting the directory by name sorts it by submission time.
    stamp = now.strftime("%Y%m%dT%H%M%S%f")[:-3] + "Z"
    payload = json.dumps(stamped, ensure_ascii=False)
    with _ANNOT_LOCK:
        # The per-submission file goes first, and it is one whole-file write:
        # that is the operation a bucket mount supports best, and it is the
        # copy the export can rebuild everything else from.
        _write_once(versions / f"{stamp}_{stamped['submission_id']}.json", payload)
        # The rolling log is a convenience, and appending to it is the part a
        # bucket mount may refuse. A failure here loses nothing, so it is
        # reported and the submission still stands.
        try:
            with (ANNOT / _ANNOT_STORES[kind][0]).open("a", encoding="utf-8") as fh:
                fh.write(payload + "\n")
                fh.flush()
                _sync(fh)
        except OSError as exc:
            print(f"annotation log append failed ({exc}); "
                  f"submission {stamped['submission_id']} kept as a file", flush=True)
    return stamped


def _sync(handle) -> None:
    """fsync where the filesystem implements it, and shrug where it does not."""
    try:
        os.fsync(handle.fileno())
    except OSError:
        pass


def _write_once(path: Path, payload: str) -> None:
    """Create a file that no later write can replace.

    Mode "x" is the guarantee; a mount that does not implement exclusive
    creation falls back to a check and a plain write, which is weaker only in
    a race that a 12-hex-character id already makes vanishingly unlikely.
    """
    try:
        with path.open("x", encoding="utf-8") as fh:
            fh.write(payload)
            fh.flush()
            _sync(fh)
    except FileExistsError:
        raise
    except OSError:
        if path.exists():
            raise FileExistsError(path)
        path.write_text(payload, encoding="utf-8")


def _stored_submissions(kind: str) -> list[dict]:
    """Every submission of this kind, oldest first.

    The rolling log is the primary source; the per-submission files fill in
    anything missing from it, so a log that was truncated, or lost with the
    container it lived in and restored from the copies, still exports whole.
    """
    rows: list[dict] = []
    seen: set[str] = set()
    log = ANNOT / _ANNOT_STORES[kind][0]
    if log.exists():
        with log.open(encoding="utf-8") as fh:
            for line in fh:
                if not line.strip():
                    continue
                try:
                    row = json.loads(line)
                except json.JSONDecodeError:
                    continue        # a half-written line, recovered below
                rows.append(row)
                if row.get("submission_id"):
                    seen.add(row["submission_id"])
    versions = ANNOT / "submissions" / kind
    if versions.is_dir():
        for path in sorted(versions.glob("*.json")):
            try:
                row = json.loads(path.read_text(encoding="utf-8"))
            except (OSError, json.JSONDecodeError):
                continue
            if row.get("submission_id") not in seen:
                rows.append(row)
    rows.sort(key=lambda r: str(r.get("submitted_at", "")))
    return rows


@app.post("/api/code/annotations")
def code_annotations(payload: dict = Body(...)) -> dict:
    records = payload.get("annotations", [])
    if not records:
        raise HTTPException(400, "No annotations")
    # `model` travels with the labels: a gold set is only comparable across models
    # if each label records the ranking it was given against.
    row = _record_submission("code", {
        "annotator": (payload.get("annotator") or "anonymous").strip(),
        "category": payload.get("category"), "query": payload.get("query"),
        "model": payload.get("model") or model_registry.DEFAULT_MODEL_ID,
        "annotations": records})
    return {"saved": len(records), "annotator": row["annotator"],
            "submission_id": row["submission_id"],
            "submitted_at": row["submitted_at"]}


# -- annotation retrieval ---------------------------------------------------
# This endpoint lets whoever runs the study pull the labels without shell
# access. On a public deployment set ENCODE_ANNOT_TOKEN so tester names and
# grades are not world-readable; when the env var is unset (local use) access
# is open.

_ANNOT_TOKEN = os.environ.get("ENCODE_ANNOT_TOKEN", "")


@app.get("/api/annotations/export")
def annotations_export(kind: str = "code", fmt: str = "csv", token: str = ""):
    if _ANNOT_TOKEN and token != _ANNOT_TOKEN:
        raise HTTPException(403, "Missing or wrong token")
    if kind not in _ANNOT_STORES:
        raise HTTPException(404, f"Unknown kind '{kind}' (use code or phenotype)")
    _, base_cols, item_cols = _ANNOT_STORES[kind]
    # Every submission ever stored, including any the rolling log lost. The
    # export is a full history, not a latest-wins view: one row per label per
    # submission, carrying the submission it belongs to.
    rows = _stored_submissions(kind)
    if not rows:
        raise HTTPException(404, f"No {kind} annotations stored yet")
    if fmt == "jsonl":
        body = "".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows)
        return StreamingResponse(
            iter([body]), media_type="application/x-ndjson",
            headers={"Content-Disposition":
                     f'attachment; filename="encode_{kind}_annotations.jsonl"',
                     "Cache-Control": _NO_CACHE})
    buf = io.StringIO()
    w = csv.writer(buf)
    w.writerow(base_cols + item_cols)
    for row in rows:
        for item in row.get("annotations", []):
            w.writerow([row.get(c, "") for c in base_cols]
                       + [item.get(c, "") for c in item_cols])
    buf.seek(0)
    return StreamingResponse(
        iter([buf.getvalue()]), media_type="text/csv",
        headers={"Content-Disposition":
                 f'attachment; filename="encode_{kind}_annotations.csv"',
                 "Cache-Control": _NO_CACHE})


# -- knowledge graph (click a code -> parent/child ontology) ---------------

@app.get("/api/graph")
def code_graph(code: str, code_type: str | None = None, drug_name: str | None = None,
               cap: int | None = None) -> dict:
    if not code.strip():
        raise HTTPException(400, "Empty code")
    return graph().neighbors(code, code_type, drug_name, cap=cap)


# -- phenotype discovery ---------------------------------------------------

def _cats(categories: str | None) -> set[str] | None:
    return {c for c in categories.split(",") if c} if categories else None


@app.get("/api/categories")
def categories() -> dict:
    return {"categories": _default_pheno_engine().categories()}


@app.get("/api/search")
def search(q: str, k: int = 10, categories: str | None = None, validated_only: bool = False,
           model: str | None = None) -> dict:
    if not q.strip():
        raise HTTPException(400, "Empty query")
    spec = _resolve_model(model, "phenotype")
    return _stamp(_pheno_engine(spec).search(q, k=_k(k),
                                             categories=_cats(categories),
                                             validated_only=validated_only), spec)


@app.get("/api/export")
def export(q: str, k: int = 10, categories: str | None = None, validated_only: bool = False,
           model: str | None = None):
    if not q.strip():
        raise HTTPException(400, "Empty query")
    spec = _resolve_model(model, "phenotype")
    data = _pheno_engine(spec).search(q, k=_k(k),
                                      categories=_cats(categories), validated_only=validated_only)
    buf = io.StringIO()
    w = csv.writer(buf)
    w.writerow(["rank", "phenotype_id", "title", "category", "validated", "relevance", "code_systems"])
    for i, r in enumerate(data["results"], 1):
        w.writerow([i, r["phenotype_id"], r["title"], r["category"], r["validated"],
                    r["scores"]["relevance"], "; ".join(r["code_systems"])])
    buf.seek(0)
    fname = f"encode_phenotype_{q.strip().replace(' ', '_')[:30]}.csv"
    return StreamingResponse(iter([buf.getvalue()]), media_type="text/csv",
                             headers={"Content-Disposition": f'attachment; filename="{fname}"'})


@app.get("/api/phenotype/{pid}")
def phenotype(pid: int) -> dict:
    detail = _default_pheno_engine().phenotype(pid)
    if detail is None:
        raise HTTPException(404, "Phenotype not found")
    return detail


@app.get("/api/phenotype/{pid}/graph")
def phenotype_graph(pid: int, focus: str | None = None, cap: int | None = None) -> dict:
    g = _default_pheno_engine().phenotype_code_graph(pid, focus=focus, cap=cap)
    if g is None:
        raise HTTPException(404, "Phenotype not found")
    return g


@app.post("/api/annotations")
def annotations(payload: dict = Body(...)) -> dict:
    """Persist phenotype-level relevance labels (the Part A evaluation gold set)."""
    records = payload.get("annotations", [])
    if not records:
        raise HTTPException(400, "No annotations")
    row = _record_submission("phenotype", {
        "annotator": (payload.get("annotator") or "anonymous").strip(),
        "query": payload.get("query"),
        "model": payload.get("model") or model_registry.DEFAULT_MODEL_ID,
        "annotations": records})
    return {"saved": len(records), "annotator": row["annotator"],
            "submission_id": row["submission_id"],
            "submitted_at": row["submitted_at"]}


@app.get("/")
def index() -> FileResponse:
    return FileResponse(FRONTEND / "index.html", headers={"Cache-Control": _NO_CACHE})


app.mount("/", NoCacheStaticFiles(directory=FRONTEND), name="static")