"""`parking-trainer serve` — the job API behind the collector's Training section.

A tiny stdlib HTTP server (no framework, no extra deps) on the compose-internal network,
never published: the collector proxies to it behind the reviewer's login. One job at a
time; each job is the CLI run as a SUBPROCESS (`python -m trainer.cli …`) with its output
captured to a log file — torch's memory goes away with the process, and a crashing job
cannot take the service down. Job state + logs persist under `<out>/jobs/` so a restart
still shows history.

  GET  /health                      {ok, busy, version}
  GET  /readiness                   what `inspect` prints (+ the defaults the UI offers)
  GET  /versions                    every <out>/<version>/ folder: written?, metrics, sidecar
  GET  /versions/<v>/report         report.md (text/markdown)
  GET  /jobs                        recent jobs, newest first
  GET  /jobs/<id>                   one job incl. the log tail
  POST /jobs                        {kind: train|evaluate|publish, …args} → 202 {id} | 409 busy
"""

from __future__ import annotations

import json
import os
import re
import subprocess
import sys
import threading
import time
import uuid
from datetime import datetime, timezone
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any

from .cli import METRICS_FILE, MODEL_FILE, REPORT_FILE, SIDECAR_FILE
from .data import load_labelled, make_split, summarise
from .preprocess import Sidecar

_VERSION_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$")
BACKBONES = ("resnet18", "mobilenet_v3_small", "efficientnet_b0")
MODES = ("features", "finetune")
LOG_TAIL_BYTES = 16_000


def _now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


class Jobs:
    """The single-slot job runner. `start` refuses while one runs."""

    def __init__(self, data_dir: Path, out_dir: Path, publish_url: str, publish_token: str) -> None:
        self.data_dir = data_dir
        self.out_dir = out_dir
        self.publish_url = publish_url
        self.publish_token = publish_token
        self.jobs_dir = out_dir / "jobs"
        self.jobs_dir.mkdir(parents=True, exist_ok=True)
        self._lock = threading.Lock()
        self._current: dict[str, Any] | None = None
        self._proc: subprocess.Popen[bytes] | None = None

    # ---- state -----------------------------------------------------------------------
    def _write(self, job: dict[str, Any]) -> None:
        (self.jobs_dir / f"{job['id']}.json").write_text(json.dumps(job, indent=2))

    def _read(self, job_id: str) -> dict[str, Any] | None:
        p = self.jobs_dir / f"{job_id}.json"
        if not p.is_file():
            return None
        return json.loads(p.read_text())  # type: ignore[no-any-return]

    def log_tail(self, job_id: str) -> str:
        p = self.jobs_dir / f"{job_id}.log"
        if not p.is_file():
            return ""
        size = p.stat().st_size
        with p.open("rb") as f:
            if size > LOG_TAIL_BYTES:
                f.seek(size - LOG_TAIL_BYTES)
            return f.read().decode("utf-8", "replace")

    @property
    def busy(self) -> bool:
        return self._current is not None

    def current(self) -> dict[str, Any] | None:
        return dict(self._current) if self._current else None

    def get(self, job_id: str) -> dict[str, Any] | None:
        if self._current and self._current["id"] == job_id:
            job = dict(self._current)
        else:
            job = self._read(job_id) or {}
        if not job:
            return None
        job["log"] = self.log_tail(job_id)
        return job

    def recent(self, limit: int = 20) -> list[dict[str, Any]]:
        files = sorted(self.jobs_dir.glob("*.json"), key=lambda p: p.stat().st_mtime, reverse=True)
        out = []
        for p in files[:limit]:
            try:
                out.append(json.loads(p.read_text()))
            except ValueError:
                continue
        if self._current and all(j["id"] != self._current["id"] for j in out):
            out.insert(0, dict(self._current))
        out.sort(key=lambda j: j.get("startedAt", ""), reverse=True)
        return out

    # ---- args → CLI ------------------------------------------------------------------
    def _argv(self, kind: str, a: dict[str, Any]) -> list[str]:
        base = [sys.executable, "-m", "trainer.cli"]
        if kind == "train":
            mode = a.get("mode", "features")
            backbone = a.get("backbone", "resnet18")
            if mode not in MODES or backbone not in BACKBONES:
                raise ValueError("bad mode/backbone")
            floor = float(a.get("minAccuracy", 0.85))
            min_per = int(a.get("minPerClass", 20))
            epochs = int(a.get("epochs", 0))
            if not 0.0 <= floor <= 1.0 or min_per < 1 or epochs < 0:
                raise ValueError("bad numbers")
            argv = base + [
                "train",
                "--data",
                str(self.data_dir),
                "--out",
                str(self.out_dir),
                "--mode",
                mode,
                "--backbone",
                backbone,
                "--min-accuracy",
                str(floor),
                "--min-per-class",
                str(min_per),
                "--epochs",
                str(epochs),
            ]
            if a.get("version"):
                argv += ["--version", self._version(a["version"])]
            return argv
        if kind == "evaluate":
            v = self._version(a.get("version", ""))
            return base + [
                "evaluate",
                "--data",
                str(self.data_dir),
                "--model",
                str(self.out_dir / v / MODEL_FILE),
            ]
        if kind == "publish":
            v = self._version(a.get("version", ""))
            url = str(a.get("url") or self.publish_url)
            if not url.startswith("https://") and not url.startswith("http://"):
                raise ValueError("bad publish url")
            return base + ["publish", str(self.out_dir / v), "--url", url]
        raise ValueError("kind must be train, evaluate or publish")

    @staticmethod
    def _version(v: Any) -> str:
        if not isinstance(v, str) or not _VERSION_RE.match(v) or v in ("cache", "jobs"):
            raise ValueError("bad version")
        return v

    # ---- run -------------------------------------------------------------------------
    def start(self, kind: str, args: dict[str, Any]) -> dict[str, Any]:
        argv = self._argv(kind, args)
        with self._lock:
            if self._current is not None:
                raise RuntimeError("busy")
            job_id = f"{datetime.now(timezone.utc).strftime('%Y%m%d-%H%M%S')}-{uuid.uuid4().hex[:6]}"
            job: dict[str, Any] = {
                "id": job_id,
                "kind": kind,
                "args": {k: v for k, v in args.items() if k != "token"},
                "status": "running",
                "startedAt": _now(),
                "finishedAt": None,
                "exitCode": None,
            }
            env = dict(os.environ)
            if kind == "publish":
                env["TRAINER_PUBLISH_TOKEN"] = self.publish_token
            log = (self.jobs_dir / f"{job_id}.log").open("wb")
            log.write(f"$ {' '.join(argv[3:])}\n".encode())
            log.flush()
            self._proc = subprocess.Popen(argv, stdout=log, stderr=subprocess.STDOUT, env=env)
            self._current = job
            self._write(job)
        threading.Thread(target=self._wait, args=(job, log), daemon=True).start()
        return dict(job)

    def _wait(self, job: dict[str, Any], log: Any) -> None:
        assert self._proc is not None
        code = self._proc.wait()
        log.close()
        with self._lock:
            job["exitCode"] = code
            job["finishedAt"] = _now()
            job["status"] = "done" if code == 0 else ("refused" if code in (2, 3) else "failed")
            self._write(job)
            self._current = None
            self._proc = None


# ----------------------------------------------------------------------------------


def readiness(data_dir: Path, min_per_class: int = 20, val_fraction: float = 0.2) -> dict[str, Any]:
    try:
        samples, missing = load_labelled(data_dir)
    except FileNotFoundError:
        return {
            "ready": False,
            "labelled": {"total": 0, "byClass": {}},
            "missingCrops": 0,
            "error": "no collector database yet",
        }
    split = make_split(samples, val_fraction, min_per_class)
    return {
        "ready": len(split.classes) >= 2,
        "labelled": summarise(samples),
        "missingCrops": missing,
        "run": {
            "classes": list(split.classes),
            "train": split.counts("train"),
            "val": split.counts("val"),
            "dropped": split.dropped,
            "minPerClass": min_per_class,
            "valFraction": val_fraction,
        },
        "defaults": {
            "mode": "features",
            "backbone": "resnet18",
            "minAccuracy": 0.85,
            "minPerClass": min_per_class,
        },
        "modes": list(MODES),
        "backbones": list(BACKBONES),
    }


def versions(out_dir: Path) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    if not out_dir.is_dir():
        return out
    for d in sorted(out_dir.iterdir(), key=lambda p: p.name, reverse=True):
        if not d.is_dir() or d.name in ("cache", "jobs"):
            continue
        if not (d / REPORT_FILE).is_file() and not (d / METRICS_FILE).is_file():
            continue
        entry: dict[str, Any] = {
            "version": d.name,
            "written": (d / MODEL_FILE).is_file() and (d / SIDECAR_FILE).is_file(),
            "hasReport": (d / REPORT_FILE).is_file(),
            "modifiedAt": datetime.fromtimestamp(d.stat().st_mtime, tz=timezone.utc)
            .replace(microsecond=0)
            .isoformat(),
        }
        try:
            m = json.loads((d / METRICS_FILE).read_text())
            entry["accuracy"] = m.get("accuracy")
            entry["macroRecall"] = m.get("macro_recall")
            entry["n"] = m.get("n")
        except (OSError, ValueError):
            pass
        if entry["written"]:
            try:
                side = Sidecar.read(d / SIDECAR_FILE)
                entry["classes"] = side.classes
                entry["mode"] = side.mode
                entry["backbone"] = side.backbone
                entry["trainedAt"] = side.trained_at
                entry["labels"] = side.labels
                entry["floor"] = side.metrics.get("floor")
            except (OSError, ValueError):
                pass
        out.append(entry)
    return out


# ----------------------------------------------------------------------------------


class Handler(BaseHTTPRequestHandler):
    jobs: Jobs  # set on the class by serve()
    server_version = "parking-trainer"

    def log_message(self, fmt: str, *args: Any) -> None:  # quieter than the default
        sys.stderr.write(f"[trainer.serve] {self.address_string()} {fmt % args}\n")

    def _json(self, code: int, body: Any) -> None:
        raw = json.dumps(body).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(raw)))
        self.end_headers()
        self.wfile.write(raw)

    def _text(self, code: int, body: str, ctype: str = "text/markdown; charset=utf-8") -> None:
        raw = body.encode()
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(raw)))
        self.end_headers()
        self.wfile.write(raw)

    def do_GET(self) -> None:  # noqa: N802
        path = self.path.split("?", 1)[0]
        j = self.jobs
        if path == "/health":
            self._json(200, {"ok": True, "busy": j.busy, "version": "parking-trainer"})
        elif path == "/readiness":
            self._json(200, readiness(j.data_dir))
        elif path == "/versions":
            self._json(200, {"versions": versions(j.out_dir)})
        elif path.startswith("/versions/") and path.endswith("/report"):
            v = path[len("/versions/") : -len("/report")]
            try:
                p = j.out_dir / Jobs._version(v) / REPORT_FILE
            except ValueError:
                self._json(400, {"error": "bad version"})
                return
            if not p.is_file():
                self._json(404, {"error": "no report"})
            else:
                self._text(200, p.read_text())
        elif path == "/jobs":
            self._json(200, {"jobs": j.recent(), "current": j.current()})
        elif path.startswith("/jobs/"):
            job = j.get(path[len("/jobs/") :])
            self._json(200, job) if job else self._json(404, {"error": "no such job"})
        else:
            self._json(404, {"error": "not found"})

    def do_POST(self) -> None:  # noqa: N802
        if self.path.split("?", 1)[0] != "/jobs":
            self._json(404, {"error": "not found"})
            return
        n = int(self.headers.get("Content-Length") or 0)
        if n > 64_000:
            self._json(413, {"error": "too large"})
            return
        try:
            body = json.loads(self.rfile.read(n) or b"{}")
            if not isinstance(body, dict):
                raise ValueError("object expected")
        except ValueError as exc:
            self._json(400, {"error": f"bad json: {exc}"})
            return
        kind = str(body.pop("kind", ""))
        try:
            job = self.jobs.start(kind, body)
        except ValueError as exc:
            self._json(400, {"error": str(exc)})
            return
        except RuntimeError:
            self._json(409, {"error": "a job is already running", "current": self.jobs.current()})
            return
        self._json(202, job)


def serve(data_dir: Path, out_dir: Path, host: str, port: int, publish_url: str, publish_token: str) -> None:
    Handler.jobs = Jobs(data_dir, out_dir, publish_url, publish_token)
    httpd = ThreadingHTTPServer((host, port), Handler)
    sys.stderr.write(f"[trainer.serve] listening on {host}:{port}, data {data_dir}, out {out_dir}\n")
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        httpd.server_close()


def wait_idle(jobs: Jobs, timeout: float = 60.0) -> None:
    """Test helper: block until no job runs."""
    t0 = time.monotonic()
    while jobs.busy and time.monotonic() - t0 < timeout:
        time.sleep(0.1)