"""LiveCodeBench grading with the semantics of the official harness (see lcb_official_runner.py), next to our own
lcb_tests.py. run_official(code, item, tests) -> (passed, total, first failure description), stops at the first failure.
As a script: lcb_official.py <tag> [items.jsonl]  regrades results/<tag>/*.jsonl with both harnesses, writes
results/<tag>/graded_official.jsonl and prints both accuracies and the disagreements."""
import glob
import json
import os
import subprocess
import sys
import tempfile
from concurrent.futures import ThreadPoolExecutor

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import sandbox  # noqa: E402

RUNNER = os.path.join(os.path.dirname(os.path.abspath(__file__)), "lcb_official_runner.py")
TMP = os.path.join(tempfile.gettempdir(), "lfm3_lcb_off")
os.makedirs(TMP, exist_ok=True)


def run_official(code, item, tests, timeout=6):
    payload = {"code": code, "mode": "functional" if item["testtype"] == "functional" else "stdin",
               "fn": item.get("func_name"), "tests": tests, "timeout": timeout}
    fd, path = tempfile.mkstemp(suffix=".json", dir=TMP)
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        json.dump(payload, f)
    try:
        r = sandbox.run([sys.executable, "-I", RUNNER, path], input=b"", capture_output=True,
                        timeout=min(30 + timeout * len(tests), 900), cwd=TMP,
                        creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
        lines = [json.loads(l) for l in r.stdout.decode(errors="replace").splitlines() if l.startswith("{")]
    except subprocess.TimeoutExpired:
        lines = []
    finally:
        try:
            os.remove(path)
        except OSError:
            pass
    ok = sum(1 for d in lines if d["res"] == "ok")
    bad = next((d for d in lines if d["res"] != "ok"), None)
    if bad is None and ok == len(tests):
        return ok, len(tests), None
    if bad is None:
        return ok, len(tests), "the run ended without a result (crash or time limit)"
    return ok, len(tests), f"test {bad['k']}: {bad['res']} {bad.get('err') or bad.get('got') or ''}"[:500]


if __name__ == "__main__":
    from grade import extract_code, final_answer
    from lcb_tests import run_lcb
    tag = sys.argv[1]
    items_path = sys.argv[2] if len(sys.argv) > 2 else "data/lcb/lcb6.jsonl"
    items = {json.loads(l)["id"]: json.loads(l) for l in open(items_path, encoding="utf-8")}
    recs = []
    for p in sorted(glob.glob(f"results/{tag}/*.jsonl")):
        if "graded" not in os.path.basename(p):
            recs += [json.loads(l) for l in open(p, encoding="utf-8")]
    recs = [r for r in recs if r["id"] in items]

    def one(r):
        it = items[r["id"]]
        code = extract_code(final_answer(r["completion"]) or "")
        if not code:
            return r, False, False, "no code"
        k, n, msg = run_official(code, it, it["hidden"])
        k2, n2, _ = run_lcb(code, it, it["hidden"])
        return r, k == n, k2 == n2, msg

    with ThreadPoolExecutor(8) as ex:
        rows = list(ex.map(one, recs))
    with open(f"results/{tag}/graded_official.jsonl", "w", encoding="utf-8") as f:
        for r, off, ours, msg in rows:
            f.write(json.dumps({"id": r["id"], "sample": r.get("sample", 0), "correct": off, "ours": ours,
                                "msg": msg}, ensure_ascii=False) + "\n")
    n = len(rows)
    print(f"{tag}: n={n}  ours {100 * sum(x[2] for x in rows) / n:.1f}  official {100 * sum(x[1] for x in rows) / n:.1f}"
          f"  (official-only {sum(x[1] and not x[2] for x in rows)}, ours-only {sum(x[2] and not x[1] for x in rows)})")
    for r, off, ours, msg in rows:
        if off != ours:
            print(f"   {r['id']} s{r.get('sample', 0)}: ours {ours} official {off} | {msg if not off else ''}"[:300])