"""Vast.ai Serverless handler: same intake + vendored chain, FastAPI surface.

Not wired to traffic (INFERENCE_BACKEND stays zerogpu). When activation is
scheduled: build the image below, create the Serverless endpoint from
vast-config.json, then flip Pages envs — frontend needs zero code change.

Response shape matches frontend/src/lib/radar.ts vastai branch:
{"scores": {finding: score}, "rows": [{finding, score}]} sorted desc.
Error strings mirror intake.py/app.py; surfaced as HTTP 400 {"detail": msg}.
"""

import os
import re
import shutil
import sys
import tempfile
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, UploadFile  # noqa: E402
import gradio as gr  # noqa: E402 — reused for intake's error literals
from huggingface_hub import snapshot_download  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))

def _vendor_path() -> str:
    """Locate vendored RADAR_inference in repo checkout or image layout (/srv)."""
    ancestors = [HERE]
    parent = HERE
    for _ in range(3):
        parent = os.path.dirname(parent)
        ancestors.append(parent)
    for base in ancestors + ["/srv", os.getcwd()]:
        if os.path.isdir(os.path.join(base, "RADAR_inference")):
            return os.path.join(base, "RADAR_inference")
    return os.path.join(HERE, "..", "..", "..", "RADAR_inference")

sys.path.insert(0, _vendor_path())
from intake import (  # noqa: E402
    HEADER_SUFFIXES,
    VOLUME_SUFFIXES,
    _series_to_nifti,
    _stage_uploads,
    _to_nifti,
)
from inference_demo import evaluate, initialize  # noqa: E402

REPO_ID = "radar-generalist/RADAR"
CHECKPOINT_NAME = "checkpoint_radar_pretrain.pth"

_state: dict = {}


def _english_name(column: str) -> str:
    # Same mapping as app.py; duplicated (not imported) so this service never
    # imports the Space module, whose import has CUDA/GPU side effects.
    m = re.search(r"\((.+)\)\s*$", column)
    return m.group(1) if m else column


def _ensure_ckpts() -> None:
    # Mirrors app.py _ensure_ckpts patterns: checkpoint + tokenizer both land
    # in MODEL_ROOT. infer_text_embedding_radar.pt ships in ckpt/ (git LFS —
    # run `git lfs pull` before docker build or it bakes the pointer).
    model_root = os.getenv("MODEL_ROOT", "/srv/ckpt")
    snapshot_download(
        REPO_ID,
        allow_patterns=[CHECKPOINT_NAME, "bert-base-chinese/*"],
        local_dir=model_root,
    )


@asynccontextmanager
async def lifespan(app: FastAPI):
    _ensure_ckpts()
    pad_func, model = initialize()  # once; per-request work stays in /score
    _state["pad_func"] = pad_func
    _state["model"] = model
    yield
    _state.clear()


app = FastAPI(title="RADAR Vast.ai scoring", lifespan=lifespan)


@app.post("/score")
async def score(ct_files: list[UploadFile]):
    """Stage uploads (intake.py) then run the vendored chain once."""
    tmpdir = tempfile.mkdtemp(prefix="radar_case_")
    outdir = tempfile.mkdtemp(prefix="radar_out_")
    incoming = tempfile.mkdtemp(prefix="radar_incoming_")
    staged = tempfile.mkdtemp(prefix="radar_dcm_")
    try:
        try:
            # Raw uploads land in incoming/; _stage_uploads alone populates
            # staged/ (it also sanitizes names and expands zips).
            paths: list[str] = []
            for i, upload in enumerate(ct_files or []):
                raw = ((upload.filename or "").strip() or f"upload_{i}").replace("\\", "/")
                safe = os.path.basename(raw) or f"upload_{i}"
                dest = os.path.join(incoming, f"{i}_{safe}")
                with open(dest, "wb") as dst:
                    dst.write(await upload.read())
                paths.append(dest)

            singles = [p for p in paths if p.lower().endswith(VOLUME_SUFFIXES)]
            rest = [p for p in paths if p not in singles]
            if singles and rest:
                raise gr.Error("upload either one volume file or one DICOM series, not both")
            if len(singles) > 1:
                raise gr.Error("upload a single volume file")
            if singles:
                src = _to_nifti(singles[0], tmpdir)
                fname = os.path.basename(src)
                if src != os.path.join(tmpdir, fname):
                    if not fname.endswith((".nii", ".nii.gz")):
                        fname += ".nii.gz"
                    shutil.copy(src, os.path.join(tmpdir, fname))
            else:
                if len(paths) == 1 and paths[0].lower().endswith(".dcm"):
                    raise gr.Error(
                        "a single .dcm is one slice, not a volume: upload the full series"
                    )
                headers = [p for p in paths if p.lower().endswith(HEADER_SUFFIXES)]
                if headers:
                    raise gr.Error(
                        f"{os.path.basename(headers[0])} needs its raw pair: "
                        "convert to .nii.gz/.nrrd/.mha first"
                    )
                _stage_uploads(paths, staged)
                _series_to_nifti(staged, tmpdir)

            import pandas as pd

            try:
                evaluate(_state["pad_func"], _state["model"], tmpdir, outdir, "vast")
            except (OSError, ValueError, RuntimeError) as exc:
                raise gr.Error(f"could not process this volume: {exc}")
            import glob

            matches = glob.glob(os.path.join(outdir, "RADAR_infer_results_vast.csv"))
            if not matches:
                raise gr.Error("model skipped this volume (check dimensions/spacing)")
            df = pd.read_csv(matches[0], encoding="utf-8-sig")
            if df.empty:
                raise gr.Error("model skipped this volume (check dimensions/spacing)")
            row = df.iloc[0]
            scores = {
                _english_name(col): float(row[col])
                for col in df.columns[1:]
                if pd.notna(row[col]) and str(row[col]).strip() != ""
            }
            rows = [
                {"finding": k, "score": v}
                for k, v in sorted(scores.items(), key=lambda kv: kv[1], reverse=True)
            ]
            return {"scores": scores, "rows": rows}
        except gr.Error as exc:
            raise HTTPException(status_code=400, detail=str(exc))
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)
        shutil.rmtree(outdir, ignore_errors=True)
        shutil.rmtree(incoming, ignore_errors=True)
        shutil.rmtree(staged, ignore_errors=True)