#!/usr/bin/env python3
"""Run ledger: one directory per training run under ``runs/``, tracked in git.

Every run is ``runs/<id>/run.json`` with a status that moves through fixed states;
each change appends to ``history`` and is meant to be committed, so
``git log runs/<id>/run.json`` is the timeline of that run. ``metrics.json`` next
to it holds the gold measure and the gate verdict, ``train.log`` the trainer's
output.

    python tools/ledger.py new --dataset dataset/2026-10-02_alc --tag alc --base yolo26n-seg.pt
    python tools/ledger.py set 2026-10-02_alc training --kv host=mx-spark-01 job=12345
    python tools/ledger.py metrics 2026-10-02_alc path/to/metrics.json
    python tools/ledger.py show
    python tools/ledger.py render        # runs/README.md
    python tools/ledger.py check         # CI: every run.json valid
"""

from __future__ import annotations

import argparse
import getpass
import json
import shutil
import socket
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

ROOT = Path(__file__).resolve().parents[1]
STATES = [
    "exported",     # dataset tag exists, run directory created
    "queued",       # job submitted (Spark or cluster)
    "training",     # epochs running, train.log grows
    "measured",     # gold measure done, metrics.json written
    "gate_passed",  # measure beats the deployed model's baseline
    "gate_failed",
    "promoted",     # models/active/production.json points at this run's weights
    "deployed",     # AareLC_infer tagged, cps-deployments bumped
    "aborted",
]
REQUIRED = ("id", "dataset", "base_weights", "status", "created", "history")


def _now() -> str:
    return datetime.now().isoformat(timespec="seconds")


def _who() -> str:
    try:
        return f"{getpass.getuser()}@{socket.gethostname()}"
    except Exception:
        return "unknown"


def runs_dir(root: Path = ROOT) -> Path:
    return root / "runs"


def run_path(run_id: str, root: Path = ROOT) -> Path:
    return runs_dir(root) / run_id / "run.json"


def load(run_id: str, root: Path = ROOT) -> Dict[str, Any]:
    p = run_path(run_id, root)
    if not p.exists():
        raise SystemExit(f"no such run: {run_id} ({p})")
    return json.loads(p.read_text(encoding="utf-8"))


def save(run: Dict[str, Any], root: Path = ROOT) -> Path:
    p = run_path(run["id"], root)
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps(run, indent=2, sort_keys=False) + "\n", encoding="utf-8")
    return p


def _parse_kv(items: Optional[List[str]]) -> Dict[str, Any]:
    out: Dict[str, Any] = {}
    for it in items or []:
        if "=" not in it:
            raise SystemExit(f"expected key=value, got {it!r}")
        k, v = it.split("=", 1)
        try:
            out[k] = json.loads(v)
        except json.JSONDecodeError:
            out[k] = v
    return out


def new_run(dataset: str, base_weights: str, tag: str, *, note: str = "", root: Path = ROOT,
            extra: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    stamp = datetime.now().strftime("%Y-%m-%d")
    run_id = f"{stamp}_{tag}"
    n = 2
    while run_path(run_id, root).exists():
        run_id = f"{stamp}_{tag}-{n}"
        n += 1
    run = {
        "id": run_id,
        "dataset": dataset,
        "base_weights": base_weights,
        "status": "exported",
        "created": _now(),
        "created_by": _who(),
        "hyperparams": {},
        "version": None,          # models/registry/<version> once trained
        "metrics": {},            # summary of metrics.json
        "history": [{"status": "exported", "at": _now(), "by": _who(), "note": note}],
    }
    run.update(extra or {})
    save(run, root)
    return run


def set_status(run_id: str, status: str, *, note: str = "", kv: Optional[Dict[str, Any]] = None,
               root: Path = ROOT) -> Dict[str, Any]:
    if status not in STATES:
        raise SystemExit(f"unknown state {status!r}; one of {', '.join(STATES)}")
    run = load(run_id, root)
    run["status"] = status
    for k, v in (kv or {}).items():
        run[k] = v
    run["history"].append({"status": status, "at": _now(), "by": _who(), "note": note})
    save(run, root)
    return run


def record_metrics(run_id: str, metrics_file: Path, root: Path = ROOT) -> Dict[str, Any]:
    run = load(run_id, root)
    data = json.loads(Path(metrics_file).read_text(encoding="utf-8"))
    dest = run_path(run_id, root).parent / "metrics.json"
    if Path(metrics_file).resolve() != dest.resolve():
        shutil.copy2(metrics_file, dest)
    overall = data.get("overall") or {}
    gate = data.get("gate") or {}
    run["metrics"] = {
        "f1": overall.get("f1"),
        "mean_mask_iou": overall.get("mean_mask_iou"),
        "gate_pass": gate.get("pass"),
        "baseline": gate.get("baseline"),
        "measured_at": _now(),
    }
    save(run, root)
    return run


def all_runs(root: Path = ROOT) -> List[Dict[str, Any]]:
    out = []
    for p in sorted(runs_dir(root).glob("*/run.json")):
        try:
            out.append(json.loads(p.read_text(encoding="utf-8")))
        except json.JSONDecodeError as e:
            out.append({"id": p.parent.name, "status": f"INVALID JSON: {e}", "history": []})
    return out


def check(root: Path = ROOT) -> List[str]:
    problems: List[str] = []
    for p in sorted(runs_dir(root).glob("*/run.json")):
        try:
            run = json.loads(p.read_text(encoding="utf-8"))
        except json.JSONDecodeError as e:
            problems.append(f"{p}: invalid JSON ({e})")
            continue
        for k in REQUIRED:
            if k not in run:
                problems.append(f"{p}: missing {k}")
        if run.get("id") != p.parent.name:
            problems.append(f"{p}: id {run.get('id')!r} does not match directory")
        if run.get("status") not in STATES:
            problems.append(f"{p}: unknown status {run.get('status')!r}")
        for h in run.get("history", []):
            if h.get("status") not in STATES:
                problems.append(f"{p}: history entry with unknown status {h.get('status')!r}")
        if run.get("status") in ("measured", "gate_passed", "gate_failed", "promoted", "deployed"):
            if not (p.parent / "metrics.json").exists():
                problems.append(f"{p}: status {run['status']} but no metrics.json")
    return problems


def _fmt(v: Any) -> str:
    if v is None:
        return "-"
    if isinstance(v, float):
        return f"{v:.3f}"
    if isinstance(v, bool):
        return "yes" if v else "no"
    return str(v)


def table(root: Path = ROOT) -> str:
    rows = [("run", "status", "dataset", "base", "version", "f1", "mask IoU", "gate", "last change")]
    for r in all_runs(root):
        m = r.get("metrics") or {}
        last = (r.get("history") or [{}])[-1].get("at", "")
        rows.append((r.get("id", "?"), r.get("status", "?"), r.get("dataset", "-"), r.get("base_weights", "-"),
                     _fmt(r.get("version")), _fmt(m.get("f1")), _fmt(m.get("mean_mask_iou")),
                     _fmt(m.get("gate_pass")), last[:16]))
    widths = [max(len(str(row[i])) for row in rows) for i in range(len(rows[0]))]
    lines = []
    for n, row in enumerate(rows):
        lines.append("  ".join(str(c).ljust(widths[i]) for i, c in enumerate(row)))
        if n == 0:
            lines.append("  ".join("-" * w for w in widths))
    return "\n".join(lines)


def render(root: Path = ROOT) -> Path:
    runs = all_runs(root)
    md = ["# Training runs", "", "Generated by `python tools/ledger.py render`; do not edit by hand.", "",
          "| Run | Status | Dataset | Base | Version | F1 | Mask IoU | Gate | Last change |",
          "| --- | --- | --- | --- | --- | --- | --- | --- | --- |"]
    for r in runs:
        m = r.get("metrics") or {}
        last = (r.get("history") or [{}])[-1].get("at", "")
        md.append(f"| {r.get('id')} | {r.get('status')} | {r.get('dataset', '-')} | {r.get('base_weights', '-')} | "
                  f"{_fmt(r.get('version'))} | {_fmt(m.get('f1'))} | {_fmt(m.get('mean_mask_iou'))} | "
                  f"{_fmt(m.get('gate_pass'))} | {last[:16]} |")
    md.append("")
    md.append("States: " + " → ".join(STATES[:-1]) + f" (or {STATES[-1]}).")
    md.append("")
    for r in runs:
        md.append(f"## {r.get('id')}")
        md.append("")
        for h in r.get("history", []):
            note = f" — {h['note']}" if h.get("note") else ""
            md.append(f"- {h.get('at', '')[:19]} · {h.get('status')} · {h.get('by', '')}{note}")
        md.append("")
    out = runs_dir(root) / "README.md"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text("\n".join(md), encoding="utf-8")
    return out


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--root", default=str(ROOT), help="repository root (default: this checkout)")
    sub = ap.add_subparsers(dest="cmd", required=True)
    n = sub.add_parser("new", help="create a run in state 'exported'")
    n.add_argument("--dataset", required=True, help="dataset tag, e.g. dataset/2026-10-02_alc")
    n.add_argument("--tag", required=True, help="short name; run id = <date>_<tag>")
    n.add_argument("--base", default="yolo26n-seg.pt", help="base weights the run starts from")
    n.add_argument("--note", default="")
    n.add_argument("--kv", nargs="*", help="extra fields key=value (JSON values accepted)")
    s = sub.add_parser("set", help="move a run to a state")
    s.add_argument("run_id")
    s.add_argument("status", choices=STATES)
    s.add_argument("--note", default="")
    s.add_argument("--kv", nargs="*", help="extra fields key=value")
    m = sub.add_parser("metrics", help="record a metrics.json on a run")
    m.add_argument("run_id")
    m.add_argument("metrics_file")
    sub.add_parser("show", help="print the table")
    sub.add_parser("render", help="write runs/README.md")
    sub.add_parser("check", help="validate every run.json (exit 1 on problems)")
    args = ap.parse_args(argv)
    root = Path(args.root)

    if args.cmd == "new":
        run = new_run(args.dataset, args.base, args.tag, note=args.note, root=root, extra=_parse_kv(args.kv))
        print(run["id"])
    elif args.cmd == "set":
        run = set_status(args.run_id, args.status, note=args.note, kv=_parse_kv(args.kv), root=root)
        print(f"{run['id']}: {run['status']}")
    elif args.cmd == "metrics":
        run = record_metrics(args.run_id, Path(args.metrics_file), root=root)
        print(json.dumps(run["metrics"]))
    elif args.cmd == "show":
        print(table(root))
    elif args.cmd == "render":
        print(render(root))
    elif args.cmd == "check":
        problems = check(root)
        for p in problems:
            print(p, file=sys.stderr)
        print(f"{len(list(runs_dir(root).glob('*/run.json')))} runs, {len(problems)} problems")
        return 1 if problems else 0
    return 0


if __name__ == "__main__":
    raise SystemExit(main())