#!/usr/bin/env python3
"""Grade a prediction file under ONE oracle, with the harness self-tests run first.

    ./grade.py --bench <bench.jsonl> --gen <gen.jsonl> --out <report.json>          # our oracle
    ./grade.py --bench <bench.jsonl> --gen <gen.jsonl> --oracle theirs --out ...    # PyLingual's
    ./grade.py --bench <bench.jsonl> --self-test-only --out preflight.json

`--oracle ours` needs nothing but the Python standard library. `--oracle theirs` needs the
optional, user-installed PyLingual extra and exits with a clear message if it is absent.

Pre-flight and the mutation test run before any score is computed. If either is not 100% the
command REFUSES to print a score.
"""
from __future__ import annotations

import argparse
import json
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from common import (  # noqa: E402
    load_bench, load_jsonl, ours_ok, pylingual_available, require_sound, self_test, strip_fences,
    theirs_ok,
)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--bench", required=True)
    ap.add_argument("--gen")
    ap.add_argument("--out", required=True)
    ap.add_argument("--oracle", choices=("ours", "theirs"), default="ours")
    ap.add_argument("--label", default="")
    ap.add_argument("--self-test-only", action="store_true")
    a = ap.parse_args()

    if a.oracle == "theirs" and not pylingual_available():
        raise SystemExit(
            "--oracle theirs needs the optional PyLingual extra, which is not installed.\n"
            "It is GPL-3.0 and is never vendored here; install it yourself (harness/README.md),\n"
            "or use --oracle ours, which reproduces every number of ours without it."
        )

    bench, _ = load_bench(a.bench)
    print(f"self-testing the {a.oracle} oracle on {len(bench)} labels...", file=sys.stderr, flush=True)
    st = self_test(bench, a.oracle)
    print(f"  PRE-FLIGHT {st['preflight_perfect']}/{st['preflight_n']} = {st['preflight_pct']}%   "
          f"MUTATION {st['mutation_killed']}/{st['mutation_total']} killed = "
          f"{st['mutation_kill_rate_pct']}%", file=sys.stderr, flush=True)
    require_sound(st)

    rep = {"label": a.label, "bench": str(a.bench), **st}
    Path(a.out).parent.mkdir(parents=True, exist_ok=True)

    if a.self_test_only or not a.gen:
        Path(a.out).write_text(json.dumps(rep, indent=2))
        print(json.dumps({k: v for k, v in rep.items()
                          if k not in ("preflight_failures", "mutation_survivors")}, indent=2))
        return

    gen = load_jsonl(a.gen)
    n = perfect = 0
    why_counts: dict[str, int] = {}
    rows = []
    with tempfile.TemporaryDirectory() as td:
        tmp = Path(td)
        for g in gen:
            i = g["i"]
            if i not in bench:
                continue
            n += 1
            src = strip_fences(g["got"])
            if a.oracle == "ours":
                ok = ours_ok(src, bench[i]["expected"])
                why = "PERFECT" if ok else "not byte-identical"
            else:
                ok, why = theirs_ok(src, Path(bench[i]["pyc_path"]), tmp, f"g{i}")
            perfect += ok
            if not ok:
                key = why.split(":")[0]
                why_counts[key] = why_counts.get(key, 0) + 1
            rows.append({"i": i, "perfect": ok, "why": why,
                         "func": bench[i].get("csn_func", ""), "n_instr": bench[i].get("n_instr")})

    rep.update({
        "gen": str(a.gen),
        "scored_n": n,
        "PERFECT": perfect,
        "PERFECT_pct": round(100 * perfect / max(1, n), 2),
        "failure_profile": dict(sorted(why_counts.items(), key=lambda x: -x[1])),
    })
    Path(a.out).write_text(json.dumps({**rep, "rows": rows}, indent=2))
    print(json.dumps({k: v for k, v in rep.items()
                      if k not in ("preflight_failures", "mutation_survivors")}, indent=2))


if __name__ == "__main__":
    main()