"""FastAPI viewer for Harbor train + eval runs. The data lives in a bucket; this only renders it.

ONE READ PATH, TWO DEPLOYMENTS. A Hugging Face bucket mounted into a Space appears as an ordinary
filesystem path inside the container (`Volume(type="bucket", source=..., mount_path="/data")`), so the
Space and a laptop run the *same* code against the same layout — `DATA_DIR` is just `/data` in one case
and `./data` in the other. The alternative, branching between a local reader and an `HfFileSystem`
reader, means two code paths where only one is ever exercised by whoever is debugging.

`DATA_BUCKET` exists for the case the mount cannot cover: reading a remote bucket from a laptop without
mounting it. It goes through `HfFileSystem`, which speaks `hf://buckets/<ns>/<name>/<path>`. When both
are set the local directory wins, because a mount is always fresher and cheaper than the network.

    DATA_DIR     = ./data                 # a local dir, or a mounted bucket at /data
    DATA_BUCKET  = (unset)                # e.g. AdithyaSK/harbor-runs — remote fallback
    REFRESH_TTL  = 30                     # seconds a listing is trusted before rescanning

WHY A REFRESH BUTTON AND NOT A WATCHER. A run in progress appends to `metrics.jsonl` and drops new files
into `traces/`; a mounted bucket reflects that without the app doing anything. But listings are cached
so that a page load does not restat thousands of files, so the UI needs a way to say "look again now" —
`POST /api/refresh` drops every cache. Nothing is precomputed, so refresh is the only invalidation
needed.

HIERARCHY: project -> dataset -> task (row) -> cell (model x harness) -> attempts (pass@k) -> trace.
Cells are keyed "<model>|<harness>" so the viewer can pivot which axis is the column without refetching;
that is a view choice, and storing it per-axis would force a rewrite to flip the table.

Endpoints:
    GET  /                                          → the viewer
    GET  /api/projects                              → tree: projects, their datasets and runs
    GET  /api/projects/{pid}/datasets/{did}         → the table (tasks + cells + aggregates)
    GET  /api/projects/{pid}/datasets/{did}/trace?path=  → one attempt
    GET  /api/projects/{pid}/runs/{rid}             → a training run's run.json
    GET  /api/projects/{pid}/runs/{rid}/train       → its metrics.jsonl, parsed
    POST /api/refresh                               → drop caches; returns what it now sees
    GET  /healthz                                   → {ok, source, n_projects, ...}
"""

from __future__ import annotations

import json
import os
import time
from pathlib import PurePosixPath
from typing import Any

from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import HTMLResponse, JSONResponse

HERE = os.path.dirname(os.path.abspath(__file__))
SITE = os.path.join(HERE, "site")
DATA_DIR = os.getenv("DATA_DIR", os.path.join(HERE, "data"))
DATA_BUCKET = os.getenv("DATA_BUCKET", "")
REFRESH_TTL = float(os.getenv("REFRESH_TTL", "30"))

app = FastAPI(title="Harbor run viewer", docs_url="/docs")

_cache: dict[str, Any] = {}
_cache_at: dict[str, float] = {}


# --- storage: a local tree, or a bucket over fsspec -------------------------------------------------
class _Local:
    kind = "local"

    def __init__(self, root: str) -> None:
        self.root = root

    def describe(self) -> str:
        return f"local:{self.root}"

    def exists(self, rel: str) -> bool:
        return os.path.exists(os.path.join(self.root, rel))

    def listdir(self, rel: str) -> list[str]:
        p = os.path.join(self.root, rel)
        return sorted(os.listdir(p)) if os.path.isdir(p) else []

    def read_text(self, rel: str) -> str:
        with open(os.path.join(self.root, rel), encoding="utf-8") as f:
            return f.read()

    def mtime(self, rel: str) -> float:
        try:
            return os.path.getmtime(os.path.join(self.root, rel))
        except OSError:
            return 0.0


class _Bucket:
    """`hf://buckets/<ns>/<name>/<path>` over HfFileSystem. Used only when nothing is mounted."""

    kind = "bucket"

    def __init__(self, bucket: str) -> None:
        from huggingface_hub import HfFileSystem

        self.bucket = bucket
        self.fs = HfFileSystem()

    def _p(self, rel: str) -> str:
        return f"buckets/{self.bucket}/{rel}".rstrip("/")

    def describe(self) -> str:
        return f"hf://buckets/{self.bucket}"

    def exists(self, rel: str) -> bool:
        return bool(self.fs.exists(self._p(rel)))

    def listdir(self, rel: str) -> list[str]:
        try:
            return sorted(PurePosixPath(p).name for p in self.fs.ls(self._p(rel), detail=False))
        except FileNotFoundError:
            return []

    def read_text(self, rel: str) -> str:
        with self.fs.open(self._p(rel), "r") as f:
            return f.read()

    def mtime(self, rel: str) -> float:
        return 0.0  # not exposed uniformly; the refresh button is the invalidation story


def store():
    """Prefer a real directory: a mounted bucket is fresher and cheaper than the network."""
    if os.path.isdir(DATA_DIR):
        return _Local(DATA_DIR)
    if DATA_BUCKET:
        return _Bucket(DATA_BUCKET)
    return _Local(DATA_DIR)  # missing; every read 404s with a path in the message


def _cached(key: str, produce):
    now = time.time()
    if key in _cache and now - _cache_at.get(key, 0.0) < REFRESH_TTL:
        return _cache[key]
    value = produce()
    _cache[key] = value
    _cache_at[key] = now
    return value


def _safe_rel(*parts: str) -> str:
    """Join a relative path and refuse to escape the data directory.

    Trace paths come straight from summary.json, which this app did not write, so a `../` in one of them
    must not read outside DATA_DIR.
    """
    for part in parts:
        if not part or part.startswith("/") or ".." in PurePosixPath(part).parts:
            raise HTTPException(400, f"unsafe path component: {part!r}")
    return str(PurePosixPath(*parts))


def _read_json(st, rel: str) -> Any:
    if not st.exists(rel):
        raise HTTPException(404, f"missing: {st.describe()}/{rel}")
    try:
        return json.loads(st.read_text(rel))
    except json.JSONDecodeError as exc:
        raise HTTPException(500, f"{rel} is not valid JSON: {exc}") from exc


def _list_projects() -> list[dict]:
    """A directory with project.json is a project. Datasets and runs are listed alongside it so the
    sidebar can render the whole tree from one request instead of N+1."""
    st = store()
    out = []
    for pid in st.listdir("projects"):
        rel = f"projects/{pid}/project.json"
        if not st.exists(rel):
            continue
        try:
            meta = json.loads(st.read_text(rel))
        except Exception:  # noqa: BLE001 - one malformed project must not hide the others
            meta = {"project_id": pid, "error": "project.json is unreadable"}
        meta.setdefault("project_id", pid)
        meta["datasets"] = [
            d for d in st.listdir(f"projects/{pid}/datasets")
            if st.exists(f"projects/{pid}/datasets/{d}/summary.json")
        ]
        meta["runs"] = [
            r for r in st.listdir(f"projects/{pid}/runs")
            if st.exists(f"projects/{pid}/runs/{r}/run.json")
        ]
        out.append(meta)
    out.sort(key=lambda p: p.get("label") or p["project_id"])
    return out


@app.get("/", response_class=HTMLResponse)
def index() -> HTMLResponse:
    p = os.path.join(SITE, "viewer.html")
    if not os.path.exists(p):
        raise HTTPException(503, "site/viewer.html is missing")
    with open(p, encoding="utf-8") as f:
        return HTMLResponse(f.read())


@app.get("/api/projects")
def projects() -> JSONResponse:
    return JSONResponse(_cached("projects", _list_projects))


@app.get("/api/projects/{pid}/datasets/{did}")
def dataset(pid: str, did: str) -> JSONResponse:
    """The table: tasks as rows, cells keyed '<model>|<harness>'. The viewer pivots which axis is the
    column, so this is returned once and re-rendered client-side rather than fetched per view."""
    rel = _safe_rel("projects", pid, "datasets", did, "summary.json")
    return JSONResponse(_cached(f"ds:{pid}/{did}", lambda: _read_json(store(), rel)))


@app.get("/api/projects/{pid}/datasets/{did}/trace")
def trace(pid: str, did: str, path: str = Query(..., description="path relative to the dataset dir")) -> JSONResponse:
    """One attempt. `path` comes from summary.json, which this app did not write, so it is validated
    against escaping the data directory before being opened."""
    rel = _safe_rel("projects", pid, "datasets", did, *PurePosixPath(path).parts)
    return JSONResponse(_read_json(store(), rel))


@app.get("/api/projects/{pid}/runs/{rid}")
def run(pid: str, rid: str) -> JSONResponse:
    return JSONResponse(
        _cached(f"run:{pid}/{rid}", lambda: _read_json(store(), _safe_rel("projects", pid, "runs", rid, "run.json")))
    )


@app.get("/api/projects/{pid}/runs/{rid}/train")
def run_train(pid: str, rid: str) -> JSONResponse:
    """metrics.jsonl -> list. A malformed final line is skipped rather than fatal: a run being appended
    to right now can have a half-written line, and refusing the file would make live runs unviewable."""

    def produce():
        st = store()
        rel = _safe_rel("projects", pid, "runs", rid, "train", "metrics.jsonl")
        if not st.exists(rel):
            raise HTTPException(404, f"missing: {st.describe()}/{rel}")
        rows, skipped = [], 0
        for line in st.read_text(rel).splitlines():
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError:
                skipped += 1
        return {"steps": rows, "skipped_lines": skipped}

    return JSONResponse(_cached(f"train:{pid}/{rid}", produce))


@app.post("/api/refresh")
def refresh() -> JSONResponse:
    """Drop every cache and rescan. A mounted bucket already reflects new writes; this is what makes
    the UI notice them without a restart."""
    _cache.clear()
    _cache_at.clear()
    listing = _list_projects()
    return JSONResponse({
        "refreshed": True,
        "source": store().describe(),
        "projects": len(listing),
        "datasets": sum(len(p.get("datasets", [])) for p in listing),
        "runs": sum(len(p.get("runs", [])) for p in listing),
    })


@app.get("/healthz")
def healthz() -> dict:
    st = store()
    return {
        "ok": os.path.exists(os.path.join(SITE, "viewer.html")),
        "source": st.describe(),
        "source_kind": st.kind,
        "data_dir_exists": os.path.isdir(DATA_DIR),
        "bucket_fallback": DATA_BUCKET or None,
        "n_projects": len(_list_projects()),
        "refresh_ttl_s": REFRESH_TTL,
    }