"""Sequential-batch eval harness for a single-GPU box: batches of samples go teacher (concurrent
vLLM) -> king reads -> challenger reads -> compare, keeping states files within the TTL and the
shm budget. Same scoring as the eval service.

    .venv/bin/python eval_seq.py --dataset pick20.jsonl --teacher DIR --king DIR --challenger DIR
        [--samples N] [--batch 8] [--workers 8] [--out results.jsonl]
"""
from __future__ import annotations

import argparse
import json
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from eval_service import dataset, metrics, run, servers
from eval_service.run import Distance, _render, generate, load_tokenizer


def aggregate(sums: list, role: str) -> Distance:
    rows = [by_role[role] for by_role in sums if role in by_role]
    if not rows:
        return Distance(kl=float("inf"), dist=float("inf"))
    return Distance(
        kl=sum(s.kl for s in rows) / sum(s.predictions for s in rows),
        dist=sum(s.layer_mean_dist for s in rows) / max(sum(s.activations for s in rows), 1),
    )


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--dataset", type=Path, required=True)
    ap.add_argument("--teacher", type=Path, required=True)
    ap.add_argument("--king", type=Path, required=True)
    ap.add_argument("--challenger", type=Path, required=True)
    ap.add_argument("--samples", type=int, default=20)
    ap.add_argument("--batch", type=int, default=8)
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--out", type=Path, default=Path("eval_seq.jsonl"))
    args = ap.parse_args()

    samples = dataset.load(args.dataset).samples[: args.samples]
    tokenizer = load_tokenizer(args.teacher)
    models = {"teacher": args.teacher, "king": args.king, "challenger": args.challenger}
    device = "cuda:0"
    heads = metrics.heads(models, [device])
    all_sums = []
    records_out = []

    try:
        for b0 in range(0, len(samples), args.batch):
            batch = samples[b0 : b0 + args.batch]
            records = []
            groups = servers.serve({"teacher": args.teacher})
            with ThreadPoolExecutor(args.workers) as pool:
                def one(i, sample, prompt):
                    answer, tpath = generate(groups["teacher"], prompt, i, sample.id)
                    for _ in range(120):  # the connector writes states asynchronously
                        if Path(tpath).exists():
                            break
                        time.sleep(0.5)
                    keep = Path("/root/states") / Path(tpath).name
                    keep.write_bytes(Path(tpath).read_bytes())
                    Path(str(keep) + ".lock").touch()
                    return answer, str(keep)

                futures = {i: (s, p, pool.submit(one, b0 + i, s, p)) for i, (s, p) in enumerate((sample, _render(tokenizer, sample.messages, sample.tools)) for sample in batch)}
                for i, (sample, prompt, fut) in futures.items():
                    answer, tpath = fut.result()
                    records.append({"id": sample.id, "prompt": prompt, "answer": answer, "tpath": tpath})
                    print(f"teacher {b0+i+1}/{len(samples)}: {sample.id} p{len(prompt)} a{len(answer)}", flush=True)
            servers.release("teacher")

            for role, dir_ in (("king", args.king), ("challenger", args.challenger)):
                groups = servers.serve({role: dir_})
                group = groups[role]
                for i, rec in enumerate(records):
                    read, path = group.complete(rec["prompt"], len(rec["answer"]), b0 + i, rec["prompt"] + rec["answer"])
                    if read != rec["answer"]:
                        raise SystemExit(f"{role} left the teacher's answer on {rec['id']}")
                    rec[role] = path
                    print(f"{role} {b0+i+1}/{len(samples)}: {rec['id']}", flush=True)
                servers.release(role)

            for rec in records:
                read = rec["prompt"] + rec["answer"][:-1]
                states = {
                    role: metrics.load_states(rec[role] if role != "teacher" else rec["tpath"], read, len(rec["prompt"]) - 1)
                    for role in ("teacher", "king", "challenger")
                }
                by_role = metrics.compare(states, heads[device], rec["answer"], device)
                all_sums.append(by_role)
                records_out.append({"id": rec["id"], **{role: s.summary() for role, s in by_role.items()}})
                print(
                    f"scored {rec['id']}: king KL {by_role['king'].kl / by_role['king'].predictions:.4f}"
                    f" challenger KL {by_role['challenger'].kl / by_role['challenger'].predictions:.4f}",
                    flush=True,
                )
                for p in (rec["tpath"], rec["king"], rec["challenger"]):
                    servers.discard(p)

        agg = {role: aggregate(all_sums, role) for role in ("teacher", "king", "challenger")}
        for role, d in agg.items():
            print(f"{role:12s} KL {d.kl:.6f}  distance {d.dist:.5f}")
        if agg["king"].kl < float("inf"):
            need = Distance(kl=0.015 * agg["king"].kl, dist=0.0084 * agg["king"].dist)
            print(
                f"margin needed: KL {need.kl:.6f}  dist {need.dist:.6f} ;"
                f" challenger leads by KL {agg['king'].kl - agg['challenger'].kl:.6f}"
                f" ({100 * (agg['king'].kl - agg['challenger'].kl) / agg['king'].kl:.3f}%)"
                f" dist {agg['king'].dist - agg['challenger'].dist:.6f}"
                f" ({100 * (agg['king'].dist - agg['challenger'].dist) / agg['king'].dist:.3f}%)"
            )
        args.out.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in records_out))
    finally:
        servers.stop_all()


if __name__ == "__main__":
    main()