"""Verify Beam's coding answers against TACO's tests on Modal (modal_app/code_verify.py). Nothing runs locally.

  python3 scripts/verify_code.py data/code/beam/run-code-*.jsonl --out data/code/verdicts.jsonl
  python3 scripts/verify_code.py --reference 300 --out data/code/reference_check.jsonl   # checks the checker

--reference runs TACO's own verified solutions for a random sample of problems; nearly all should pass.
"""

import argparse
import glob
import json
import random
from pathlib import Path

import modal

ROOT = Path(__file__).resolve().parent.parent
CHUNK = 25


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("files", nargs="*")
    ap.add_argument("--out", required=True)
    ap.add_argument("--reference", type=int, default=0)
    ap.add_argument("--skip", nargs="*", default=[], help="verdict files whose answers are already verified")
    args = ap.parse_args()

    tests = {t["id"]: t for t in map(json.loads, open(ROOT / "data" / "code" / "tests.jsonl"))}
    done = {json.loads(l)["key"] for f in args.skip for l in open(f)}
    records = []
    if args.reference:
        raw = {f"taco-{r['id']}": r for r in json.load(open(ROOT / "data" / "code_raw" / "taco" / "taco_verified.json"))}
        for pid in random.Random(0).sample(sorted(tests), args.reference):
            sols = raw[pid]["solutions"]
            sol = (json.loads(sols) if isinstance(sols, str) else sols)[0]
            records.append({"key": pid, "content": f"```python\n{sol}\n```", "test": tests[pid]})
    for pattern in args.files:
        for path in glob.glob(pattern):
            for line in open(path):
                r = json.loads(line)
                if "error" in r or r["finish_reason"] != "stop" or r["id"] not in tests:
                    continue
                if f"{r['id']}#{r.get('sample_idx', 0)}" in done:
                    continue
                records.append({"key": f"{r['id']}#{r.get('sample_idx', 0)}", "content": r["content"],
                                "test": tests[r["id"]]})
    print(f"{len(records)} answers to verify")

    fn = modal.Function.from_name("beam-dory-codeverify", "verify", environment_name="main")
    chunks = [records[i:i + CHUNK] for i in range(0, len(records), CHUNK)]
    n = passed = 0
    with open(args.out, "w") as f:
        for res in fn.map(chunks, order_outputs=False):
            for v in res:
                f.write(json.dumps(v) + "\n")
                n += 1
                passed += v["passed"]
            if n % 1000 < CHUNK:
                print(f"  {n}/{len(records)} verified, {passed} passed", flush=True)
    print(f"done: {passed}/{n} passed ({100 * passed / max(n, 1):.1f}%) -> {args.out}")


if __name__ == "__main__":
    main()