#!/usr/bin/env python3
"""Aggregate 42-task benchmark results into results.json + results_full42.csv.

Source of truth: the private `brywil/gguf-quant-lab` checkout's
`benchmarks/results_<label>_<suite>.json` files.

Scoring contract
----------------
- A task is **passed** only when every pytest item in it passed.
- A **void** is a harness limit (timeout, output-cap truncation, no completion).
  Voids are excluded from capability: `valid = tasks - void`.
- The headline is `passed / valid`, not `passed / tasks`. Never count a void as a failure.

Usage
-----
    python scripts/aggregate_results.py --lab ~/gguf-quant-lab --out .

Outputs
-------
    results.json          all full-42 and subset runs, with per-suite/total fields
    results_full42.csv    flat one-row-per-run table

A run is "full-42" when all four suites are present. Labels are grouped by stripping a
trailing `-rN` / `-repN` so repeated runs of the same arm compare together.
"""
from __future__ import annotations

import argparse
import csv
import glob
import json
import os
import re
from collections import defaultdict


SUITES = [
    ("opencode_tasks", "general", 25),
    ("opencode_tasks_hard", "hard", 5),
    ("opencode_tasks_frontier", "frontier", 9),
    ("opencode_tasks_tier4a", "tier4a", 3),
]


def model_path(lab: str, label: str) -> str:
    """Best-effort: the served GGUF basename from a local server log, if retained."""
    for path in sorted(glob.glob(os.path.join(lab, "benchmarks", "results", f"{label}*server.log"))):
        try:
            text = open(path, errors="ignore").read(200_000)
        except OSError:
            continue
        m = re.search(r"loading model '([^']+)'", text)
        if m:
            return os.path.basename(m.group(1))
    return ""


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--lab", default=os.path.expanduser("~/gguf-quant-lab"),
                    help="checkout containing benchmarks/results_*.json")
    ap.add_argument("--out", default=".", help="output directory")
    args = ap.parse_args()

    per_label: dict[str, dict] = defaultdict(dict)
    for path in glob.glob(os.path.join(args.lab, "benchmarks", "results_*.json")):
        base = os.path.basename(path)[len("results_"):-len(".json")]
        for suffix, name, _ in SUITES:
            if not base.endswith("_" + suffix):
                continue
            label = base[: -(len(suffix) + 1)]
            try:
                data = json.load(open(path))
            except (OSError, json.JSONDecodeError):
                continue
            challenges = data.get("challenges", []) or []
            tasks = len(challenges)
            void = int(data.get("n_void", 0) or 0)
            passed = sum(1 for c in challenges if c.get("passed"))
            partial = (sum(float(c.get("score", 0) or 0) for c in challenges) / tasks) if tasks else 0.0
            per_label[label][name] = {
                "tasks": tasks,
                "valid": tasks - void,
                "passed": passed,
                "void": void,
                "partial": round(partial, 4),
                "wall_s": round(float(data.get("wall_s", 0) or 0), 1),
                "tokens_out": sum(int(c.get("tokens_out", 0) or 0) for c in challenges),
                "source": os.path.basename(path),
            }
            break

    runs, partial_runs = [], []
    for label, suites in sorted(per_label.items()):
        record = {
            "label": label,
            "group": re.sub(r"-(r|rep)\d+$", "", label),
            "box": "thor" if (
                label.lower().startswith("thor-") or
                label.lower().startswith("thor_") or
                label.lower().endswith("_thor") or
                label.lower().endswith("-thor") or
                "_thor_" in label.lower() or
                "-thor-" in label.lower() or
                "_thor-" in label.lower()
            ) else "gb10",
            "model_path": model_path(args.lab, label),
            "suites": suites,
            "total_passed": sum(s["passed"] for s in suites.values()),
            "total_valid": sum(s["valid"] for s in suites.values()),
            "total_void": sum(s["void"] for s in suites.values()),
            "total_wall_s": round(sum(s["wall_s"] for s in suites.values()), 1),
            "total_tokens_out": sum(s["tokens_out"] for s in suites.values()),
        }
        (runs if len(suites) == len(SUITES) else partial_runs).append(record)

    os.makedirs(args.out, exist_ok=True)
    payload = {
        "generated_from": "brywil/gguf-quant-lab benchmarks/results_*.json",
        "scoring": "passed/valid; valid = tasks - void; void = timeout/truncation/no completion (not a failure)",
        "suites": {name: {"dir": suffix, "tasks": count} for suffix, name, count in SUITES},
        "runs": runs,
        "partial_runs": partial_runs,
    }
    with open(os.path.join(args.out, "results.json"), "w") as f:
        json.dump(payload, f, indent=2)

    csv_path = os.path.join(args.out, "results_full42.csv")
    with open(csv_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["label", "group", "box", "model_path",
                         "general", "hard", "frontier", "tier4a",
                         "valid_total", "void", "passed_total",
                         "wall_h", "tokens_out", "tokens_per_task", "tok_s", "partial_mean"])
        for r in sorted(runs, key=lambda x: (x["group"], x["label"])):
            s = r["suites"]
            writer.writerow([
                r["label"], r["group"], r["box"], r["model_path"],
                f"{s['general']['passed']}/{s['general']['valid']}",
                f"{s['hard']['passed']}/{s['hard']['valid']}",
                f"{s['frontier']['passed']}/{s['frontier']['valid']}",
                f"{s['tier4a']['passed']}/{s['tier4a']['valid']}",
                f"{r['total_passed']}/{r['total_valid']}", r["total_void"],
                r["total_passed"], round(r["total_wall_s"] / 3600, 2),
                r["total_tokens_out"],
                round(r["total_tokens_out"] / max(1, r["total_valid"]), 1),
                round(r["total_tokens_out"] / max(1, r["total_wall_s"]), 2),
                round(sum(s[k]["partial"] * s[k]["tasks"] for k in s) / max(1, r["total_valid"]), 4),
            ])

    print(f"full-42 runs: {len(runs)} | subset runs: {len(partial_runs)}")
    print(f"wrote {os.path.join(args.out, 'results.json')} and results_full42.csv")


if __name__ == "__main__":
    main()