#!/usr/bin/env python3
"""Render the 42-task results table from models.json + results.json.

Each row is one entry in models.json. The best run for that group (highest
passed/valid, then most valid tasks, then fewest voids) supplies the numbers.

Usage
-----
    python scripts/render_table.py --models models.json --results results.json
    python scripts/render_table.py ... --readme README.md     # replace between markers

README markers (literal):
    <!-- TABLE:START -->
    ... generated table ...
    <!-- TABLE:END -->
"""
from __future__ import annotations

import argparse
import json
import re


def best_run(model: dict, runs: list[dict]) -> dict | None:
    matches = [r for r in runs if r["label"].startswith(model["group"]) and r["total_valid"] > 0]
    if not matches:
        return None
    return max(matches, key=lambda r: (r["total_passed"], r["total_valid"], -r["total_void"]))


def reps_text(model: dict, runs: list[dict]) -> str:
    matches = [r for r in runs if r["label"].startswith(model["group"])]
    if not matches:
        return "not run"
    passed = sorted({r["total_passed"] for r in matches})
    if len(matches) == 1:
        return "1"
    if len(passed) == 1:
        return f"{len(matches)} ({passed[0]})"
    return f"{len(matches)} ({passed[0]}–{passed[-1]})"


def render(models: list[dict], runs: list[dict]) -> str:
    lines = [
        "| Model | Quant | G (25) | H (5) | F (9) | T4A (3) | Valid | Void | Tok/task | Tok/s | Reps |",
        "|---|---|---:|---:|---:|---:|---:|---:|---:|---:|---|",
    ]
    for model in models:
        run = best_run(model, runs)
        name = f"[{model['model']}]({model['link']})" if model.get("link") else model["model"]
        if run is None:
            lines.append(f"| {name} | {model['quant']} | — | — | — | — | — | — | — | — | not run |")
            continue
        s = run["suites"]
        valid = max(1, run["total_valid"])
        wall = max(1, run["total_wall_s"])
        lines.append(
            f"| {name} | {model['quant']} "
            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"| **{run['total_passed']}/{run['total_valid']}** "
            f"| {run['total_void']} "
            f"| {run['total_tokens_out']/valid:,.0f} "
            f"| {run['total_tokens_out']/wall:.1f} "
            f"| {reps_text(model, runs)} |"
        )
    return "\n".join(lines)


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--models", default="models.json")
    ap.add_argument("--results", default="results.json")
    ap.add_argument("--readme", help="replace the table between TABLE:START/END markers in this file")
    args = ap.parse_args()

    models = json.load(open(args.models))["models"]
    runs = json.load(open(args.results))["runs"]
    table = render(models, runs)

    if not args.readme:
        print(table)
        return

    text = open(args.readme).read()
    pattern = re.compile(r"(<!-- TABLE:START -->\n).*?(\n<!-- TABLE:END -->)", re.S)
    if not pattern.search(text):
        raise SystemExit(f"{args.readme}: markers <!-- TABLE:START --> / <!-- TABLE:END --> not found")
    open(args.readme, "w").write(pattern.sub(lambda m: m.group(1) + table + m.group(2), text))
    print(f"updated {args.readme} ({len(models)} rows)")


if __name__ == "__main__":
    main()