"""Run one system over one or more eval sets and score it.

    CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=<GPU UUID, never GPU 0> python -m deg.bench.run --system deg --set datagen-test \
        --n 200 --out /shared/rwkv/results/deg/bench-v1

Writes ``<out>/<system>/<set>.preds.jsonl`` (one :class:`Prediction` per record; resumable:
uids already predicted without error are skipped, the last line per uid wins) and
``<out>/<system>/<set>.metrics.json`` (the aggregate of :mod:`deg.bench.score`, plus system
info and GPU utilization samples for local systems).
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any

from deg.bench import sets as bench_sets
from deg.bench.score import Prediction, aggregate, score_record
from deg.bench.systems import System, make


def load_preds(path: Path, keep_errors: bool = False) -> dict[str, Prediction]:
    out: dict[str, Prediction] = {}
    if path.exists():
        for line in path.read_text().split("\n"):
            if line.strip():
                d = json.loads(line)
                if keep_errors or d.get("error") is None:  # failures retried on resume
                    out[d["uid"]] = Prediction(**d)
    return out


def run_set(system: System, set_name: str, recs: list[dict[str, Any]], out: Path,
            log_every: int = 10) -> dict[str, Any]:
    sdir = out / bench_sets.safe_name(system.name)
    sdir.mkdir(parents=True, exist_ok=True)
    stem = bench_sets.safe_name(set_name)
    preds_path = sdir / f"{stem}.preds.jsonl"
    done = load_preds(preds_path)
    todo = [r for r in recs if r["uid"] not in done]
    monitor = None
    if system.workers == 1:
        try:
            from deg.train import GpuMonitor
            monitor = GpuMonitor()
        except Exception:  # noqa: BLE001 - telemetry only
            monitor = None
    t0 = time.time()
    with preds_path.open("a", encoding="utf-8") as fh:
        def emit(i: int, p: Prediction) -> None:
            done[p.uid] = p
            fh.write(json.dumps(p.to_json(), ensure_ascii=False) + "\n")
            fh.flush()
            if (i + 1) % log_every == 0 or i + 1 == len(todo):
                util = monitor.take() if monitor else {}
                print(json.dumps({"system": system.name, "set": set_name, "done": i + 1,
                                  "todo": len(todo), "elapsed_s": round(time.time() - t0, 1),
                                  "rec_per_s": round((i + 1) / (time.time() - t0), 3),
                                  **util}), flush=True)

        if system.workers > 1:
            with ThreadPoolExecutor(system.workers) as pool:
                for i, p in enumerate(pool.map(system.predict, todo)):
                    emit(i, p)
        else:
            for i, rec in enumerate(todo):
                emit(i, system.predict(rec))
    util_samples = monitor.samples if monitor else []
    if monitor:
        monitor.stop()
    preds = [done[r["uid"]] for r in recs]
    scored = [score_record(r, p, system.scope) for r, p in zip(recs, preds)]
    metrics = {"system": system.name, "set": set_name, "scope": list(system.scope),
               "info": system.info, **aggregate(scored, preds)}
    if util_samples:
        metrics["gpu_util_mean"] = round(sum(u for u, _ in util_samples) / len(util_samples), 1)
    (sdir / f"{stem}.metrics.json").write_text(json.dumps(metrics, indent=2))
    (sdir / f"{stem}.scored.jsonl").write_text(
        "".join(json.dumps(s, ensure_ascii=False) + "\n" for s in scored))
    return metrics


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser(description="deg head-to-head bench")
    ap.add_argument("--system", required=True,
                    choices=["deg", "constrained", "constrained-llamacpp", "nuextract", "clef",
                             "frontier"])
    ap.add_argument("--set", action="append", required=True, dest="sets")
    ap.add_argument("--n", type=int, default=None, help="records per set (manifest order)")
    ap.add_argument("--out", type=Path, default=Path("/shared/rwkv/results/deg/bench-v1"))
    ap.add_argument("--records", default=None, help="datagen records.jsonl")
    ap.add_argument("--extractbench", default=None)
    ap.add_argument("--ckpt", default=None, help="deg checkpoint or run dir (latest step)")
    ap.add_argument("--model-path", default=None, help="constrained: HF model dir")
    ap.add_argument("--compact-json", action="store_true",
                    help="constrained: ',' / ':' separators (ablation)")
    ap.add_argument("--server-url", default=None, help="nuextract/clef server")
    ap.add_argument("--models", default="gabe-codex-1/gpt-6-luna",
                    help="frontier: comma list of CPA model ids (rotated)")
    ap.add_argument("--workers", type=int, default=4)
    ap.add_argument("--label", default=None, help="system name override")
    args = ap.parse_args(argv)
    opts = {"records": args.records, "extractbench": args.extractbench}
    recs_by_set = {s: bench_sets.sample(s, args.n, args.out, opts) for s in args.sets}
    system = make(args.system, args)
    try:
        for s, recs in recs_by_set.items():
            m = run_set(system, s, recs, args.out)
            print(json.dumps({k: v for k, v in m.items() if k != "info"}), flush=True)
    finally:
        system.close()
    sys.stdout.flush()


if __name__ == "__main__":
    main()