"""
evaluate — run the three leaderboard evals for one model and publish the results.

    uv run evaluate "GPT 5.5" --model gpt-5.5 --provider openai --effort medium \\
        --provider-name OpenAI --url https://openai.com --release-date 2026-08-01 \\
        --price-input 1.25 --price-output 10
    uv run evaluate "GPT 5.5"              # resume or rerun with its recorded config
    uv run evaluate "GPT 5.5" --status     # peek at run/batch progress
    uv run evaluate --verify               # offline rescore of every published run

The first argument is the model's display name: its entry in data/models.json.
For a model already on the board, the run identity (model id, provider,
effort, ...) comes from its published result files, so a bare `evaluate
"Name"` resumes or reproduces exactly that run; any identity flag you pass
overrides it (and so starts a different run). For a new model, pass the
identity flags plus the metadata the harness cannot know (provider name, url,
release date, price).

Each eval (freeform, multiple choice, leetcode) runs as its own process, one
after another or all at once with --parallel. When all three results exist,
the entry's resultFiles (and kilowattHour, with --energy) are written to
models.json and `uv run build --write` rebuilds the site data — all
under one lock, so parallel `evaluate` processes never clobber each other.
"""

import argparse
import os
import subprocess
import sys
import threading
from pathlib import Path

from evals import luau_leetcode as lc
from evals import registry
from evals import roblox_qa as qa

# Each tag selects a model key, eval module, and optional mode.
EVALS = {
    "ff": ("freeform", qa, "--freeform"),
    "mc": ("multipleChoice", qa, "--multiple-choice"),
    "lc": ("leetcode", lc, None),
}
KEY_TO_TAG = {key: tag for tag, (key, _, _) in EVALS.items()}

# These keys name the run folder and results file. The eval sets the mode.
IDENTITY = tuple(k for k in qa.RUN_KEYS if k != "mode")

# Electricity-priced entries need these kilowattHour keys.
KWH_KEYS = ("freeform", "multipleChoice", "easy", "medium", "hard")

# Use --allow-errors above this failure limit. Dead endpoints and empty keys
# usually exceed it.
MAX_ERROR_RATE = 0.05


# --- arguments ---------------------------------------------------------------------
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    ap = argparse.ArgumentParser(
        prog="evaluate",
        description="Run the RobloxQA (freeform + multiple choice) and Luau Leetcode "
        "evals for one model, then update data/models.json and rebuild the "
        "leaderboard.",
    )
    ap.add_argument("name", nargs="?", help="display name (the models.json entry)")

    ident = ap.add_argument_group(
        "run identity",
        "Defaults come from the model's published result files. Passing one "
        "changes the run (new run folder and results file).",
    )
    ident.add_argument("--model", help="model id, or checkpoint dir for --provider local")
    ident.add_argument("--provider", choices=["live", "openai", "anthropic", "local"])
    ident.add_argument("--base-url", help="live only: OpenAI-compatible endpoint")
    reasoning = ident.add_mutually_exclusive_group()
    reasoning.add_argument("--effort", help="reasoning effort (e.g. low/medium/high)")
    reasoning.add_argument(
        "--thinking",
        action="store_true",
        default=None,
        help="anthropic yes/no extended thinking, or local enable_thinking",
    )
    ident.add_argument("--route", help="live only: pin OpenRouter provider routing")
    ident.add_argument("--max-tokens", type=int, help="candidate completion cap")
    ident.add_argument("--system-prompt-prefix", dest="system_prefix")
    ident.add_argument("--adapter", help="local only: LoRA adapter dir")
    ident.add_argument("--limit", type=int, help="smoke test: first N questions/problems")
    ident.add_argument(
        "--split",
        choices=["test", "train"],
        help="RobloxQA split; 'train' is for generation runs, never scores",
    )

    launch = ap.add_argument_group("how to run (not part of the run identity)")
    launch.add_argument(
        "--api-key-env",
        help="live only: env var holding the endpoint's key (OPENROUTER_API_KEY is "
        "picked automatically for openrouter.ai)",
    )
    launch.add_argument("--responses-api", action="store_true")
    launch.add_argument("--workers", type=int, help="live only: parallel requests per eval (default 2)")
    launch.add_argument(
        "--allow-errors",
        action="store_true",
        help=f"publish even when more than {MAX_ERROR_RATE:.0%} of a run's answers "
        "are API errors or unparseable verdicts",
    )
    launch.add_argument("--jest-url", help="JestLuaTestServer URL (default JEST_SERVER_URL)")
    launch.add_argument(
        "--evals",
        default="ff,mc,lc",
        help="comma list of evals to run: ff (freeform), mc (multiple choice), lc (leetcode). Default: all three",
    )
    launch.add_argument(
        "--parallel",
        action="store_true",
        help="run the evals at the same time instead of one after another",
    )
    launch.add_argument(
        "--force",
        action="store_true",
        help="discard existing run state and start over (re-spends money)",
    )
    launch.add_argument(
        "--energy",
        action="store_true",
        help="log request spans and write measured kilowattHour from the HWiNFO log "
        "(checks first that HWiNFO is logging)",
    )
    launch.add_argument("--energy-log", type=Path, help="HWiNFO CSV (default C:\\hwinfo\\eval.csv)")
    launch.add_argument(
        "--no-build",
        action="store_true",
        help="run the evals but leave models.json and the site data alone",
    )
    launch.add_argument("--dry-run", action="store_true", help="print the plan and exit")
    launch.add_argument("--status", action="store_true", help="print run status and exit")
    launch.add_argument(
        "--verify",
        action="store_true",
        help="rescore published runs offline and compare with their result files (all models, or just NAME)",
    )

    meta = ap.add_argument_group("models.json metadata (required for a new model, optional updates otherwise)")
    meta.add_argument("--provider-name", help='display provider, e.g. "OpenAI"')
    meta.add_argument("--url", help="model page")
    meta.add_argument("--release-date", help="YYYY-MM-DD")
    meta.add_argument(
        "--size",
        nargs="+",
        metavar="BYTES_OR_LINK",
        help="weight size: a byte count, or Hugging Face links to the weights "
        "(a file, part 1 of a split GGUF, a folder, or a repo's safetensors). "
        "Several values are added up, e.g. the GGUF and its mmproj file",
    )
    meta.add_argument("--price-input", type=float, help="USD per million input tokens")
    meta.add_argument("--price-output", type=float, help="USD per million output tokens")
    meta.add_argument("--price-electricity", type=float, help="USD per kWh (local models)")

    args = ap.parse_args(argv)
    if not args.verify and not args.name:
        ap.error("a model name is required (or pass --verify)")
    tags = [t.strip() for t in args.evals.split(",") if t.strip()]
    unknown = [t for t in tags if t not in EVALS]
    if unknown or not tags:
        ap.error(f"--evals takes a comma list of ff, mc, lc (got {args.evals!r})")
    args.tags = tags
    for v in args.size or []:
        if not (v.isdigit() or v.startswith(("http://", "https://"))):
            ap.error(f"--size takes byte counts or Hugging Face links (got {v!r})")
    if args.price_electricity is not None and (args.price_input is not None or args.price_output is not None):
        ap.error("--price-electricity excludes --price-input/--price-output")
    return args


# --- run planning --------------------------------------------------------------------
def published_configs(entry: dict | None) -> dict[str, dict]:
    """tag -> config of each published result file of this entry."""
    out = {}
    for key, fname in ((entry or {}).get("resultFiles") or {}).items():
        tag = KEY_TO_TAG.get(key)
        if tag is None or not fname:
            continue
        path = EVALS[tag][1].RESULTS_DIR / fname
        if path.exists():
            out[tag] = qa.normalize_config(qa.read_results(path)["config"])
    return out


def identity_for(tag: str, published: dict[str, dict], args) -> dict:
    """The run identity for one eval: its published config (or a sibling's),
    with any identity flag from the command line on top."""
    base = published.get(tag) or next(iter(published.values()), {})
    ident = {k: base.get(k) for k in IDENTITY if k in base}
    for k in IDENTITY:
        v = getattr(args, k)
        if v is not None:
            ident[k] = v
    return ident


def module_argv(tag: str, ident: dict, args) -> list[str]:
    """Command-line flags for one eval module."""
    _, _, mode_flag = EVALS[tag]
    argv = [mode_flag] if mode_flag else []
    provider = ident.get("provider") or "live"
    argv += ["--provider", provider]
    if ident.get("model"):
        argv += ["--model", ident["model"]]
    if provider == "live" and ident.get("base_url"):
        argv += ["--base-url", ident["base_url"]]
    for key, flag in (
        ("effort", "--effort"),
        ("route", "--route"),
        ("max_tokens", "--max-tokens"),
        ("adapter", "--adapter"),
        ("limit", "--limit"),
        ("system_prefix", "--system-prompt-prefix"),
    ):
        if ident.get(key) is not None:
            argv += [flag, str(ident[key])]
    if ident.get("thinking"):
        argv.append("--thinking")
    if tag != "lc" and ident.get("split") not in (None, "test"):
        argv += ["--split", ident["split"]]

    if args.api_key_env and provider == "live":
        argv += ["--api-key-env", args.api_key_env]
    if args.responses_api:
        argv.append("--responses-api")
    if args.workers is not None:
        argv += ["--workers", str(args.workers)]
    if tag == "lc" and args.jest_url:
        argv += ["--jest-url", args.jest_url]
    if args.energy:
        argv.append("--energy")
    if args.force:
        argv.append("--force")
    return argv


def run_for(tag: str, argv: list[str]):
    """The module's Run for these flags (its config, run folder, results path)."""
    module = EVALS[tag][1]
    run = module.Run.from_args(module.parse_args(argv))
    return run, module.results_path(run.config)


# --- running -----------------------------------------------------------------------
def _child_env() -> dict:
    env = dict(os.environ)
    env["PYTHONIOENCODING"] = "utf-8"
    env["PYTHONUNBUFFERED"] = "1"
    return env


def run_serial(jobs: list[tuple[str, list[str]]]) -> dict[str, int]:
    codes = {}
    for tag, cmd in jobs:
        print(f"\n===== {tag}: {' '.join(cmd[2:])}", flush=True)
        codes[tag] = subprocess.run(cmd, env=_child_env(), check=False).returncode
    return codes


def run_parallel(jobs: list[tuple[str, list[str]]]) -> dict[str, int]:
    """All evals at once; each output line is prefixed with its eval's tag."""
    print_lock = threading.Lock()

    def pump(tag: str, proc: subprocess.Popen) -> None:
        for line in proc.stdout:
            with print_lock:
                print(f"[{tag}] {line}", end="", flush=True)

    procs = {}
    threads = []
    for tag, cmd in jobs:
        print(f"[{tag}] {' '.join(cmd[2:])}", flush=True)
        proc = subprocess.Popen(
            cmd,
            env=_child_env(),
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            encoding="utf-8",
            errors="replace",
        )
        procs[tag] = proc
        t = threading.Thread(target=pump, args=(tag, proc), daemon=True)
        t.start()
        threads.append(t)
    codes = {tag: proc.wait() for tag, proc in procs.items()}
    for t in threads:
        t.join()
    return codes


def error_heavy(finished: dict[str, Path]) -> list[str]:
    """One line per result whose API errors and unparseable verdicts exceed
    MAX_ERROR_RATE of its answers."""
    lines = []
    for tag, path in finished.items():
        summary = qa.read_results(path)["summary"]
        summary = summary.get("overall", summary)  # leetcode nests its totals
        failed = (summary.get("apiErrors") or 0) + (summary.get("judgeUnparseable") or 0)
        if failed > MAX_ERROR_RATE * summary["n"]:
            lines.append(f"{tag}: {failed} of {summary['n']} answers failed ({path.name})")
    return lines


# --- models.json -----------------------------------------------------------------------
def resolve_size(values: list[str]) -> int:
    """--size values (byte counts and Hugging Face links) added up."""
    from evals import metadata

    total = 0
    for v in values:
        if v.isdigit():
            total += int(v)
            continue
        try:
            size = metadata.weights_size(v)
        except Exception as e:  # noqa: BLE001 — bad link, gated repo, offline
            raise SystemExit(f"error: --size {v}: {e}") from None
        print(f"  --size {v}: {size:,} bytes")
        total += size
    return total


def fill_metadata(args, config: dict) -> None:
    """Fill the metadata flags a new model left out, from OpenRouter and
    Hugging Face (evals/metadata.py), and print what was filled."""
    from evals import metadata

    found = metadata.lookup(config, args.url)
    filled = []
    for flag, attr, key in (
        ("--provider-name", "provider_name", "provider"),
        ("--url", "url", "url"),
        ("--release-date", "release_date", "releaseDate"),
        ("--size", "size", "size"),
    ):
        if getattr(args, attr) is None and key in found:
            setattr(args, attr, found[key])
            filled.append(f"{flag} {found[key]}")
    no_price = args.price_electricity is None and args.price_input is None and args.price_output is None
    if no_price and "price" in found:
        args.price_input, args.price_output = (
            found["price"]["input"],
            found["price"]["output"],
        )
        filled.append(f"--price-input {args.price_input} --price-output {args.price_output}")
    for line in filled:
        print(f"  looked up {line}")


def new_entry_fields(args, entry: dict | None) -> dict:
    """Metadata from the flags, merged over the existing entry."""
    e = dict(entry or {"name": args.name})
    if args.provider_name is not None:
        e["provider"] = args.provider_name
    if args.url is not None:
        e["url"] = args.url
    if args.release_date is not None:
        e["releaseDate"] = args.release_date
    if args.size is not None:
        e["size"] = args.size
    e.setdefault("size", None)
    if args.price_electricity is not None:
        e["price"] = {"electricity": args.price_electricity}
    elif args.price_input is not None or args.price_output is not None:
        price = dict(e.get("price") or {})
        price.pop("electricity", None)
        if args.price_input is not None:
            price["input"] = args.price_input
        if args.price_output is not None:
            price["output"] = args.price_output
        e["price"] = price
    return e


ENTRY_ORDER = (
    "name",
    "provider",
    "url",
    "releaseDate",
    "size",
    "price",
    "kilowattHour",
    "resultFiles",
)


def order_keys(e: dict) -> dict:
    """The hand-edited key order, so a new entry reads like the others."""
    return {k: e[k] for k in ENTRY_ORDER if k in e} | {k: v for k, v in e.items() if k not in ENTRY_ORDER}


def check_new_entry(args, entry: dict | None) -> None:
    if entry is not None:
        return
    missing = [
        flag
        for flag, value in (
            ("--provider-name", args.provider_name),
            ("--url", args.url),
            ("--release-date", args.release_date),
        )
        if value is None
    ]
    if args.price_electricity is None and (args.price_input is None or args.price_output is None):
        missing.append("--price-input/--price-output or --price-electricity")
    if missing:
        raise SystemExit(
            f'error: "{args.name}" is not in data/models.json yet; a new model also needs {", ".join(missing)}'
        )


def measure_energy(runs: dict[str, object], log: Path | None) -> dict[str, float]:
    """kilowattHour keys for the evals that ran, from their logged spans."""
    from evals import energy

    all_dirs = [f.parent for bench in (qa.CACHE_DIR, lc.CACHE_DIR) for f in bench.glob(f"*/{energy.INTERVALS_FILE}")]
    targets = [run.dir for run in runs.values()]
    shares = energy.attribute_runs(targets, all_dirs, log or energy.DEFAULT_LOG)
    kwh: dict[str, float] = {}
    for tag, run in runs.items():
        by_label = shares.get(run.dir, {})
        if tag == "lc":  # split by difficulty band
            for band in ("easy", "medium", "hard"):
                kwh[band] = round(by_label.get(band, 0.0), 3)
        else:
            kwh[EVALS[tag][0]] = round(by_label.get("all", 0.0), 3)
    # Keep the earlier value when this run measures zero energy.
    for k, v in kwh.items():
        if v == 0:
            print(f"  warning: no energy measured for {k} (no logged spans, or no HWiNFO samples then)")
    return {k: v for k, v in kwh.items() if v}


def publish(args, results: dict[str, Path], kwh: dict[str, float]) -> None:
    """Record finished results in models.json and rebuild the site data.

    The build accepts a model with all three resultFiles or none, and an
    electricity-priced model only with all five kilowattHour values. Until
    that holds, the entry keeps its previous (valid) state and only the
    metadata is updated."""

    def mutate(entry: dict | None) -> dict:
        e = new_entry_fields(args, entry)
        if not all(tag in results for tag in EVALS):
            have = ", ".join(t for t in EVALS if t in results) or "none"
            print(f"  finished so far: {have}; resultFiles is set once all three exist")
            return order_keys(e)
        electricity = "electricity" in (e.get("price") or {})
        kwh_all = dict(e.get("kilowattHour") or {}) | kwh
        if electricity and not all(isinstance(kwh_all.get(k), (int, float)) for k in KWH_KEYS):
            print(
                "  resultFiles not set: this model is priced by electricity and "
                "needs all five kilowattHour values (run with --energy, or add "
                "them to data/models.json by hand)"
            )
            return order_keys(e)
        if electricity:
            e["kilowattHour"] = {k: kwh_all[k] for k in KWH_KEYS}
        e["resultFiles"] = {key: results[tag].name for tag, (key, _, _) in EVALS.items()}
        return order_keys(e)

    registry.update_entry(args.name, mutate)


# --- verify --------------------------------------------------------------------------
def verify(names: list[str]) -> int:
    """Offline rescore of published runs; nonzero exit on any mismatch."""
    models = registry.load_models()
    if names:
        chosen = [m for m in models if m["name"] in names]
        unknown = set(names) - {m["name"] for m in chosen}
        if unknown:
            raise SystemExit(f"error: not in data/models.json: {sorted(unknown)}")
    else:
        chosen = [m for m in models if m.get("resultFiles")]
    bad = 0
    for m in chosen:
        for key, fname in (m.get("resultFiles") or {}).items():
            tag = KEY_TO_TAG[key]
            module = EVALS[tag][1]
            path = module.RESULTS_DIR / fname
            try:
                old = qa.read_results(path)
                summary, results = module.rescore(old["config"])
            except Exception as e:  # noqa: BLE001 — report and keep going
                print(f"FAIL  {m['name']} [{tag}] {fname}: {e}", flush=True)
                bad += 1
                continue
            diffs = []
            if summary != old["summary"]:
                diffs.append(
                    "summary: "
                    + ", ".join(
                        f"{k} {old['summary'].get(k)!r} -> {summary.get(k)!r}"
                        for k in sorted(set(summary) | set(old["summary"]))
                        if summary.get(k) != old["summary"].get(k)
                    )
                )
            rows = sum(a != b for a, b in zip(results, old["results"]))
            rows += abs(len(results) - len(old["results"]))
            if rows:
                diffs.append(f"{rows} result rows differ")
            if diffs:
                bad += 1
                print(f"DIFF  {m['name']} [{tag}] {fname}: {'; '.join(diffs)}", flush=True)
            else:
                print(f"ok    {m['name']} [{tag}] {fname}", flush=True)
    print(f"\n{'all runs reproduce' if not bad else f'{bad} run(s) do not reproduce'}")
    return 1 if bad else 0


# --- main --------------------------------------------------------------------------
def is_smoke_test(runs: dict) -> bool:
    """--limit and train-split runs never go on the board."""
    return any(
        run.config.get("limit") is not None or run.config.get("split", "test") != "test" for run in runs.values()
    )


def main(argv: list[str] | None = None) -> None:
    sys.stdout.reconfigure(line_buffering=True, encoding="utf-8", errors="replace")
    args = parse_args(argv)
    if args.verify:
        raise SystemExit(verify([args.name] if args.name else []))
    if args.size is not None:
        args.size = resolve_size(args.size)

    entry = registry.find_entry(registry.load_models(), args.name)
    if entry is None and args.model is None and args.provider is None:
        raise SystemExit("error: a new model needs --model and --provider")
    published = published_configs(entry)

    # Plan every eval. Run selected ones and use existing results for the rest.
    runs, results_paths, cmds = {}, {}, {}
    for tag in EVALS:
        argv_t = module_argv(tag, identity_for(tag, published, args), args)
        runs[tag], results_paths[tag] = run_for(tag, argv_t)
        cmds[tag] = [sys.executable, "-m", EVALS[tag][1].__name__, *argv_t]
    jobs = [(tag, cmds[tag]) for tag in args.tags]
    smoke = is_smoke_test(runs)
    if entry is None and not (args.status or smoke):
        fill_metadata(args, runs["ff"].config)
    if not (args.status or args.dry_run or args.no_build or smoke):
        check_new_entry(args, entry)  # fail before spending anything
    if args.energy and not (args.status or args.dry_run):
        from evals import energy

        energy.check_logging(args.energy_log or energy.DEFAULT_LOG)

    print(f"{args.name}:")
    for tag, cmd in jobs:
        print(f"  {tag}: {subprocess.list2cmdline(cmd[2:])}")
        print(f"      run folder {runs[tag].dir}")
        print(f"      results    {results_paths[tag]}")
    if args.dry_run:
        return
    if args.status:
        for tag in args.tags:
            qa.print_status(runs[tag])
        return

    codes = run_parallel(jobs) if args.parallel else run_serial(jobs)
    failed = [tag for tag, code in codes.items() if code != 0]
    finished = {tag: path for tag, path in results_paths.items() if tag not in failed and path.exists()}

    kwh: dict[str, float] = {}
    if args.energy and finished:
        try:
            kwh = measure_energy({t: runs[t] for t in finished}, args.energy_log)
            print(f"  measured kWh: {kwh}")
        except Exception as e:  # noqa: BLE001 — the scores still stand
            print(f"  energy attribution failed: {e}")

    heavy = error_heavy(finished)
    if smoke:
        print("  smoke test (--limit/--split): models.json left alone")
    elif args.no_build:
        print("  --no-build: models.json left alone")
    elif heavy and not args.allow_errors:
        print("  not published: too many failed answers")
        for line in heavy:
            print(f"      {line}")
        print(
            "  Rerun to retry the failed live answers. If the failures are final "
            "(batch errors), pass --allow-errors to publish anyway."
        )
    elif finished:
        publish(args, finished, kwh)

    if failed:
        raise SystemExit(f"error: {', '.join(failed)} failed (see the log above); rerun to resume")
    if heavy and not (smoke or args.no_build or args.allow_errors):
        raise SystemExit("error: not published: too many failed answers (see above)")


if __name__ == "__main__":
    main()