"""FastAPI backend for the FrameWorker user-study Space — dual-study edition.

Serves TWO independent studies from one Space:

  v1 (original)   pages at /            API at /api/*      dataset folder submissions/
  v2 (40-runs)    pages at /v2/         API at /v2/api/*   dataset folder v2/submissions/

Both write to the same dataset repo (`dehezhang2/Frameworker_User_Study`) but
into disjoint folders, so the two studies' databases never mix. The rater
pages use relative `./api/*` URLs, which resolve to the right prefix
automatically depending on which page they were loaded from.

Endpoint behavior is identical to the original single-study app.py; the
handlers are built by a factory parameterized on (prefix, submissions folder,
annotations file). The huobao→univa legacy rename only applies to the v1
study.

Secret: HF_TOKEN must be set as a Space secret with write access to the
dataset repo.
"""
from __future__ import annotations

import datetime as dt
import json
import os
import re
import uuid
from io import BytesIO

from fastapi import APIRouter, FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import HfHubHTTPError

DATASET_REPO = "dehezhang2/Frameworker_User_Study"

# Plain shared-secret gate for the admin dashboards (`/admin.html`,
# `/v2/admin.html`). Override via `ADMIN_PASSWORD` env var / Space secret.
# The default is intentionally weak — this isn't auth-grade, just a screen
# against random visitors stumbling into the submissions list.
ADMIN_PASSWORD = os.environ.get("ADMIN_PASSWORD", "12345")

# Mirror of the client-side ADMIN_NAMES list. Matched case-insensitively
# against the rater's `name` field. Server-side computation is the
# authoritative one — clients can't spoof the flag by lying in the payload.
ADMIN_NAMES = {
    "deheng zhang",
    "zhendong li",
}

app = FastAPI(title="FrameWorker user-study backend (dual-study)")

_TOKEN = os.environ.get("HF_TOKEN")
_api = HfApi(token=_TOKEN) if _TOKEN else None
_EMAIL_RE = re.compile(r"^[^@\s]+@[^@\s]+\.[^@\s]+$")


def _is_score(x) -> bool:
    return isinstance(x, int) and 1 <= x <= 5


def _slug(s: str) -> str:
    return re.sub(r"[^A-Za-z0-9._-]+", "_", s).strip("_")[:64] or "anon"


def _is_admin(name: str) -> bool:
    return (name or "").strip().lower() in ADMIN_NAMES


def _normalize_legacy_models(payload: dict) -> int:
    """v1-study only: rewrite any `model: "huobao"` (and matching
    `_huobao.mp4` paths) to `univa` on the way in. Older clients with stale
    localStorage keep re-introducing the old name on save; this silently
    fixes it server-side so the dataset stays clean."""
    n = 0
    for c in payload.get("cases") or []:
        for v in c.get("videos") or []:
            if v.get("model") == "huobao":
                v["model"] = "univa"; n += 1
            p = v.get("path") or ""
            if "_huobao.mp4" in p:
                v["path"] = p.replace("_huobao.mp4", "_univa.mp4"); n += 1
    return n


def _count_scores(cases) -> int:
    n = 0
    for c in cases or []:
        for v in c.get("videos", []) or []:
            for s in (v.get("ratings") or {}).values():
                if _is_score(s):
                    n += 1
    return n


def _validate_rater(payload) -> tuple[str, str]:
    rater = payload.get("rater") or {}
    name = (rater.get("name") or "").strip()
    email = (rater.get("email") or "").strip().lower()
    if not name:
        raise HTTPException(400, "rater.name is required")
    if not _EMAIL_RE.match(email):
        raise HTTPException(400, "rater.email must be a valid email address")
    return name, email


def _require_admin_password(req: Request) -> None:
    pw = req.headers.get("x-admin-password") or req.query_params.get("pw") or ""
    if pw != ADMIN_PASSWORD:
        raise HTTPException(401, "admin password required")


def build_study_router(
    *,
    prefix: str,
    study: str,
    submissions_root: str,
    annotations_path: str,
    normalize_legacy: bool,
    cases_file: str,
) -> APIRouter:
    """One study's full API surface, writing under its own dataset folders."""
    router = APIRouter(prefix=prefix)
    _case_ids_cache: list[str] = []

    def _manifest_case_ids() -> list[str]:
        """Case ids in manifest order, read from this study's static
        cases.json (baked into the container image — safe to cache)."""
        if not _case_ids_cache:
            try:
                with open(cases_file, "r", encoding="utf-8") as f:
                    _case_ids_cache.extend(
                        c["id"] for c in json.load(f) if c.get("id"))
            except Exception:
                pass
        return _case_ids_cache

    def _load_annotations() -> dict:
        """Return the study's current annotations file, or an empty shell."""
        if _api is None:
            return {"annotations": {}}
        try:
            local = hf_hub_download(
                repo_id=DATASET_REPO, repo_type="dataset",
                filename=annotations_path, force_download=True,
            )
            with open(local, "r", encoding="utf-8") as f:
                data = json.load(f)
            if not isinstance(data.get("annotations"), dict):
                data["annotations"] = {}
            return data
        except HfHubHTTPError as e:
            # 404 means we haven't created the file yet — that's the empty case.
            if e.response is not None and e.response.status_code == 404:
                return {"annotations": {}}
            return {"annotations": {}}
        except Exception:
            return {"annotations": {}}

    @router.get("/health")
    def health():
        return {
            "ok": True,
            "has_token": bool(_TOKEN),
            "dataset": DATASET_REPO,
            "study": study,
            "time_utc": dt.datetime.utcnow().isoformat() + "Z",
        }

    @router.post("/submit")
    async def submit(req: Request):
        if _api is None:
            raise HTTPException(500, "server missing HF_TOKEN secret")
        try:
            payload = await req.json()
        except Exception as e:
            raise HTTPException(400, f"invalid JSON: {e}")

        name, email = _validate_rater(payload)
        if normalize_legacy:
            _normalize_legacy_models(payload)

        rubric = payload.get("rubric") or []
        cases = payload.get("cases") or []
        if not rubric or not cases:
            raise HTTPException(400, "submission must include rubric + cases")

        qkeys = [q["key"] for d in rubric for q in d.get("questions", [])]
        if not qkeys:
            raise HTTPException(400, "rubric has no questions")

        missing = []
        bad = []
        n_videos = 0
        n_scores = 0
        for c in cases:
            cid = c.get("id", "?")
            for v in c.get("videos", []):
                n_videos += 1
                ratings = v.get("ratings") or {}
                slot = v.get("slot", "?")
                for qk in qkeys:
                    if qk not in ratings:
                        missing.append(f"{cid}/{slot}/{qk}")
                    elif not _is_score(ratings[qk]):
                        bad.append(f"{cid}/{slot}/{qk}={ratings[qk]!r}")
                    else:
                        n_scores += 1
        if missing or bad:
            raise HTTPException(
                400,
                json.dumps(
                    {
                        "error": "incomplete_or_invalid",
                        "missing_count": len(missing),
                        "invalid_count": len(bad),
                        "first_missing": missing[:5],
                        "first_invalid": bad[:5],
                    }
                ),
            )

        ts = dt.datetime.utcnow().strftime("%Y%m%dT%H%M%SZ")
        fname = f"{submissions_root}/{_slug(email)}_{ts}_{uuid.uuid4().hex[:8]}.json"

        is_admin = _is_admin(name)
        payload["is_admin"] = is_admin
        payload["_server"] = {
            "received_at_utc": dt.datetime.utcnow().isoformat() + "Z",
            "client_ip": req.client.host if req.client else None,
            "n_videos": n_videos,
            "n_scores": n_scores,
            "stored_as": fname,
            "status": "final",
        }
        blob = json.dumps(payload, indent=2, ensure_ascii=False).encode("utf-8")

        try:
            _api.upload_file(
                repo_id=DATASET_REPO,
                repo_type="dataset",
                path_or_fileobj=BytesIO(blob),
                path_in_repo=fname,
                commit_message=(
                    f"{study}: submission from {email} ({n_scores} scores"
                    f"{' · admin' if is_admin else ''})"
                ),
            )
        except HfHubHTTPError as e:
            status = e.response.status_code if e.response is not None else 500
            raise HTTPException(
                status_code=502,
                detail=f"failed to write to dataset ({status}): {e}",
            )

        return JSONResponse(
            {
                "ok": True,
                "stored_as": fname,
                "n_videos": n_videos,
                "n_scores": n_scores,
                "is_admin": is_admin,
            }
        )

    @router.post("/save_progress")
    async def save_progress(req: Request):
        """Write a per-rater in-progress snapshot to the dataset.

        Used in two situations:
          - Right after registration: page POSTs an empty payload (no ratings
            yet) so the rater appears in the in_progress folder from the
            moment they sign in.
          - Mid-study auto-save (debounced) + manual "Save progress" button:
            the same single per-user file is overwritten with the latest
            snapshot.

        Unlike /submit, this does NOT require completeness — the payload may
        have any subset of ratings (including none).
        """
        if _api is None:
            raise HTTPException(500, "server missing HF_TOKEN secret")
        try:
            payload = await req.json()
        except Exception as e:
            raise HTTPException(400, f"invalid JSON: {e}")

        name, email = _validate_rater(payload)
        if normalize_legacy:
            _normalize_legacy_models(payload)
        cases = payload.get("cases") or []
        n_scores = _count_scores(cases)
        is_admin = _is_admin(name)

        fname = f"{submissions_root}/in_progress/{_slug(email)}.json"
        payload["is_admin"] = is_admin
        payload["_server"] = {
            "received_at_utc": dt.datetime.utcnow().isoformat() + "Z",
            "client_ip": req.client.host if req.client else None,
            "n_scores": n_scores,
            "stored_as": fname,
            "status": "in_progress",
        }
        blob = json.dumps(payload, indent=2, ensure_ascii=False).encode("utf-8")
        try:
            _api.upload_file(
                repo_id=DATASET_REPO,
                repo_type="dataset",
                path_or_fileobj=BytesIO(blob),
                path_in_repo=fname,
                commit_message=(
                    f"{study}: in-progress save from {email} ({n_scores} scores"
                    f"{' · admin' if is_admin else ''})"
                ),
            )
        except HfHubHTTPError as e:
            status = e.response.status_code if e.response is not None else 500
            raise HTTPException(
                status_code=502,
                detail=f"failed to write to dataset ({status}): {e}",
            )

        return JSONResponse({
            "ok": True,
            "stored_as": fname,
            "n_scores": n_scores,
            "is_admin": is_admin,
        })

    @router.post("/admin/login")
    async def admin_login(req: Request):
        try:
            payload = await req.json()
        except Exception:
            raise HTTPException(400, "invalid JSON")
        if (payload.get("password") or "").strip() != ADMIN_PASSWORD:
            raise HTTPException(401, "wrong password")
        return JSONResponse({"ok": True})

    @router.get("/admin/submissions")
    async def admin_submissions(req: Request):
        """List every submission JSON in this study's folder with light
        metadata. Admin-only (X-Admin-Password header)."""
        _require_admin_password(req)
        if _api is None:
            return JSONResponse({"submissions": [], "reason": "no token"})
        out = []
        try:
            tree_iter = _api.list_repo_tree(
                repo_id=DATASET_REPO, repo_type="dataset",
                path_in_repo=submissions_root, recursive=True,
            )
            tree_list = list(tree_iter)
        except HfHubHTTPError as e:
            # 404 = no submissions yet for this study — empty, not an error.
            status = e.response.status_code if e.response is not None else 500
            if status == 404:
                return JSONResponse({"submissions": [], "total": 0})
            raise HTTPException(502, f"failed to list dataset ({status}): {e}")
        for tree in tree_list:
            path = getattr(tree, "path", "")
            if not path.endswith(".json"):
                continue
            record: dict = {"path": path}
            try:
                local = hf_hub_download(
                    repo_id=DATASET_REPO, repo_type="dataset",
                    filename=path, force_download=False,
                )
                with open(local, "r", encoding="utf-8") as f:
                    data = json.load(f)
                rater = data.get("rater") or {}
                srv = data.get("_server") or {}
                cases = data.get("cases") or []
                record.update({
                    "rater_name": rater.get("name") or "",
                    "rater_email": rater.get("email") or "",
                    "is_admin": bool(data.get("is_admin", False)),
                    "n_cases": len(cases),
                    "n_scores": srv.get("n_scores"),
                    "status": srv.get("status"),
                    "received_at_utc": srv.get("received_at_utc"),
                    "chosen_case_ids": rater.get("chosen_case_ids") or [],
                })
            except Exception as e:
                record["error"] = str(e)[:120]
            out.append(record)
        # newest first
        out.sort(key=lambda r: r.get("received_at_utc") or "", reverse=True)
        return JSONResponse({"submissions": out, "total": len(out)})

    @router.delete("/admin/submission")
    async def admin_delete_submission(req: Request):
        """Admin-only delete of one of this study's submission JSONs.

        Body: {"path": "<submissions_root>/...json"} — must live under this
        study's submissions folder and end with ".json" to prevent the
        endpoint from deleting unrelated dataset files via path traversal.
        """
        _require_admin_password(req)
        if _api is None:
            raise HTTPException(500, "server missing HF_TOKEN secret")
        try:
            payload = await req.json()
        except Exception:
            raise HTTPException(400, "invalid JSON")
        path = (payload.get("path") or "").strip()
        if (not path.startswith(submissions_root + "/")
                or not path.endswith(".json") or ".." in path):
            raise HTTPException(
                400, f"path must be a {submissions_root}/*.json file")
        try:
            _api.delete_file(
                repo_id=DATASET_REPO, repo_type="dataset",
                path_in_repo=path,
                commit_message=(
                    f"{study} admin: delete submission {os.path.basename(path)}"
                ),
            )
        except HfHubHTTPError as e:
            status = e.response.status_code if e.response is not None else 500
            raise HTTPException(502, f"delete failed ({status}): {e}")
        return JSONResponse({"ok": True, "deleted": path})

    @router.get("/assign_cases")
    async def assign_cases(email: str = "", n: int = 10):
        """Coverage-first case assignment for a newly registering rater.

        Returns the `n` case ids with the FEWEST distinct raters covering
        them so far (random tie-break, manifest order in the response), so
        un-rated cases get picked up before already-rated ones. A case counts
        as "covered" by a rater when that rater either has >=1 actual rating
        on it or was assigned it at registration (`chosen_case_ids`) —
        counting assignments too keeps two near-simultaneous registrants from
        landing on the identical least-rated subset. The requesting rater's
        own files are excluded.

        The client treats this as advisory: on any failure it falls back to
        its original per-rater seeded random pick.
        """
        email = (email or "").strip().lower()
        manifest_ids = _manifest_case_ids()
        if _api is None or not manifest_ids:
            return JSONResponse({"case_ids": [], "reason": "unavailable"})
        try:
            tree_list = [
                t for t in _api.list_repo_tree(
                    repo_id=DATASET_REPO, repo_type="dataset",
                    path_in_repo=submissions_root, recursive=True,
                )
                if getattr(t, "path", "").endswith(".json")
            ]
        except HfHubHTTPError as e:
            status = e.response.status_code if e.response is not None else 500
            if status != 404:  # 404 = no submissions yet — zero coverage
                return JSONResponse({"case_ids": [], "reason": f"HTTP {status}"})
            tree_list = []

        per_rater: dict[str, set] = {}
        for t in tree_list:
            try:
                local = hf_hub_download(
                    repo_id=DATASET_REPO, repo_type="dataset",
                    filename=t.path, force_download=False,
                )
                with open(local, "r", encoding="utf-8") as f:
                    data = json.load(f)
            except Exception:
                continue
            rater = data.get("rater") or {}
            em = (rater.get("email") or "").strip().lower()
            if email and em == email:
                continue
            covered = per_rater.setdefault(em or t.path, set())
            for cid in rater.get("chosen_case_ids") or []:
                covered.add(cid)
            for case in data.get("cases") or []:
                for v in case.get("videos") or []:
                    if any(_is_score(s) for s in (v.get("ratings") or {}).values()):
                        covered.add(case.get("id"))
                        break

        coverage = {cid: 0 for cid in manifest_ids}
        for covered in per_rater.values():
            for cid in covered:
                if cid in coverage:
                    coverage[cid] += 1

        import random as _random
        n = max(1, min(n, len(manifest_ids)))
        picked = sorted(
            manifest_ids,
            key=lambda c: (coverage[c], _random.random()),
        )[:n]
        order = {cid: i for i, cid in enumerate(manifest_ids)}
        picked.sort(key=lambda c: order[c])
        return JSONResponse({"case_ids": picked, "coverage": coverage})

    @router.get("/rater_state")
    async def get_rater_state(email: str = ""):
        """Return the rater's most recent in-progress JSON (if any).

        Used by the page to recover a rater's locked case subset + prior
        ratings when they land on a new device."""
        email = (email or "").strip().lower()
        if not email:
            raise HTTPException(400, "email query param required")
        slug = _slug(email)
        if _api is None:
            return JSONResponse({"found": False, "reason": "no token"})
        try:
            local = hf_hub_download(
                repo_id=DATASET_REPO, repo_type="dataset",
                filename=f"{submissions_root}/in_progress/{slug}.json",
                force_download=True,
            )
            with open(local, "r", encoding="utf-8") as f:
                return JSONResponse({"found": True, "state": json.load(f)})
        except HfHubHTTPError as e:
            status = e.response.status_code if e.response is not None else 500
            if status == 404:
                return JSONResponse({"found": False})
            return JSONResponse({"found": False, "reason": f"HTTP {status}"})
        except Exception as e:
            return JSONResponse({"found": False, "reason": str(e)})

    @router.get("/annotations")
    async def get_annotations():
        """Public read of admin-curated prompt highlights.

        Shape: {"annotations": {case_id: {lang: [snippet, snippet, ...]}}}
        """
        data = _load_annotations()
        return JSONResponse({"annotations": data.get("annotations", {})})

    @router.post("/annotations")
    async def post_annotations(req: Request):
        """Admin-only: replace the highlight list for a (case_id, lang) pair.

        Payload: {rater:{name,email}, case_id, lang, ranges: [string, ...]}
        Server-enforces ADMIN_NAMES against the supplied rater.name.
        """
        if _api is None:
            raise HTTPException(500, "server missing HF_TOKEN secret")
        try:
            payload = await req.json()
        except Exception as e:
            raise HTTPException(400, f"invalid JSON: {e}")
        name, email = _validate_rater(payload)
        if not _is_admin(name):
            raise HTTPException(403, "admin only")
        case_id = (payload.get("case_id") or "").strip()
        lang = (payload.get("lang") or "en").strip()
        ranges = payload.get("ranges") or []
        if not case_id:
            raise HTTPException(400, "case_id required")
        if not isinstance(ranges, list) or not all(isinstance(r, str) for r in ranges):
            raise HTTPException(400, "ranges must be a list of strings")
        # Sanitise: drop empty / huge entries, dedupe in order.
        cleaned = []
        seen = set()
        for r in ranges:
            s = r.strip()
            if not s or len(s) > 1024 or s in seen:
                continue
            seen.add(s); cleaned.append(s)

        data = _load_annotations()
        data.setdefault("annotations", {}).setdefault(case_id, {})[lang] = cleaned
        data["_server"] = {
            "last_edited_by": email,
            "last_edited_name": name,
            "last_edited_at_utc": dt.datetime.utcnow().isoformat() + "Z",
            "case_id": case_id, "lang": lang,
        }
        blob = json.dumps(data, indent=2, ensure_ascii=False).encode("utf-8")
        try:
            _api.upload_file(
                repo_id=DATASET_REPO, repo_type="dataset",
                path_or_fileobj=BytesIO(blob),
                path_in_repo=annotations_path,
                commit_message=(
                    f"{study} annotations: {email} set {len(cleaned)} "
                    f"highlight(s) for {case_id} / {lang}"
                ),
            )
        except HfHubHTTPError as e:
            status = e.response.status_code if e.response is not None else 500
            raise HTTPException(502, f"failed to write annotations ({status}): {e}")
        return JSONResponse({
            "ok": True, "case_id": case_id, "lang": lang,
            "ranges": cleaned,
        })

    return router


# v1 — the original study; paths and behavior identical to the old app.py.
app.include_router(build_study_router(
    prefix="/api", study="v1",
    submissions_root="submissions",
    annotations_path="annotations.json",
    normalize_legacy=True,
    cases_file="static/cases.json",
))

# v1 archived UI — same v1 study backend, for the pages preserved at /v1/.
app.include_router(build_study_router(
    prefix="/v1/api", study="v1",
    submissions_root="submissions",
    annotations_path="annotations.json",
    normalize_legacy=True,
    cases_file="static/v1/cases.json",
))

# v2 — the 40-runs baseline study; everything lives under v2/ in the dataset.
# The Space root (static/index.html) redirects here: v2 is the CURRENT study.
app.include_router(build_study_router(
    prefix="/v2/api", study="v2",
    submissions_root="v2/submissions",
    annotations_path="v2/annotations.json",
    normalize_legacy=False,
    cases_file="static/v2/cases.json",
))

# Static rater UIs live in ./static (v1 at /, v2 at /v2/). Mount last so the
# API routes take precedence.
app.mount("/", StaticFiles(directory="static", html=True), name="static")