"""Run Beam's coding answers against TACO's tests on Modal, with networking blocked. Generated code never runs on
the Mac.

  modal deploy modal_app/code_verify.py::APP
  python3 scripts/verify_code.py RESULTS_FILE_ON_VOLUME     # spawns verify over chunks, writes a verdict file

Each answer's last ```python block runs once per test (at most 30) in a fresh subprocess with a 6 s limit per test
and 2 GB of memory:
  stdin/stdout problems  stdout must equal the expected output, compared token by token (floats to 1e-6)
  function problems      the starter function (or Solution().<name>) must return the expected value; TACO wraps
                         most expected values in a one-element list, so both forms are accepted
"""

import json
import re
import time

import modal

APP = modal.App("beam-dory-codeverify")
image = modal.Image.debian_slim(python_version="3.11").pip_install("numpy")
data = modal.Volume.from_name("beam-dory-data", version=2)
D = "/data"

HARNESS = r'''
import json, sys
sys.setrecursionlimit(100000)
spec = json.load(open(sys.argv[1]))
ns = {"__name__": "__solution__"}
exec(compile(open(sys.argv[2]).read(), "solution.py", "exec"), ns)
if "Solution" in ns and isinstance(ns["Solution"], type):
    fn = getattr(ns["Solution"](), spec["fn_name"])
else:
    fn = ns[spec["fn_name"]]
def norm(x):
    return json.loads(json.dumps(x, default=list))
for args, want in zip(spec["inputs"], spec["outputs"]):
    got = norm(fn(*args))
    if not (got == want or [got] == want or (isinstance(want, list) and len(want) == 1 and got == want[0])):
        print("WRONG", flush=True)
        sys.exit(3)
print("OK")
'''


# LeetCode-style problems assume these imports, as LeetCode itself provides them
PREAMBLE = ("from typing import *\nimport collections, heapq, math, itertools, functools, bisect, string, re\n"
            "from collections import deque, defaultdict, Counter, OrderedDict\n"
            "from functools import lru_cache, reduce, cache\nfrom heapq import heappush, heappop, heapify\n"
            "from itertools import accumulate, permutations, combinations, product\n"
            "from bisect import bisect_left, bisect_right\n")


def extract_code(content):
    blocks = re.findall(r"```(?:python|py|python3)?[ \t]*\n(.*?)```", content or "", re.S)
    return blocks[-1] if blocks else None


def tokens_equal(got, want):
    a, b = got.split(), want.split()
    if a == b:
        return True
    if len(a) != len(b):
        return False
    for x, y in zip(a, b):
        if x == y:
            continue
        try:
            if abs(float(x) - float(y)) > 1e-6 * max(1.0, abs(float(y))):
                return False
        except ValueError:
            return False
    return True


def limits():
    import resource
    resource.setrlimit(resource.RLIMIT_AS, (2 << 30, 2 << 30))


def run_one(code, test, workdir):
    import subprocess
    from pathlib import Path

    src = Path(workdir) / "solution.py"
    if test["fn_name"]:
        src.write_text(code if code.lstrip().startswith("from __future__") else PREAMBLE + code)
        spec = Path(workdir) / "spec.json"
        spec.write_text(json.dumps(test))
        (Path(workdir) / "harness.py").write_text(HARNESS)
        try:
            p = subprocess.run(["python3", "harness.py", "spec.json", "solution.py"], cwd=workdir, capture_output=True,
                               text=True, timeout=6 * len(test["inputs"]) + 4, preexec_fn=limits)
        except subprocess.TimeoutExpired:
            return False, "timeout"
        return (p.returncode == 0 and p.stdout.strip().endswith("OK")), ("ok" if p.returncode == 0 else
                                                                         p.stdout[-80:] + p.stderr[-200:])
    src.write_text(code)
    for inp, want in zip(test["inputs"], test["outputs"]):
        inp = inp if isinstance(inp, str) else "\n".join(map(str, inp))
        want = want if isinstance(want, str) else "\n".join(map(str, want))
        try:
            p = subprocess.run(["python3", "solution.py"], cwd=workdir, input=inp, capture_output=True, text=True,
                               timeout=6, preexec_fn=limits)
        except subprocess.TimeoutExpired:
            return False, "timeout"
        if p.returncode != 0:
            return False, "error: " + p.stderr[-200:]
        if not tokens_equal(p.stdout, want):
            return False, "wrong answer"
    return True, "ok"


@APP.function(image=image, volumes={D: data}, cpu=2, memory=8192, timeout=1800, block_network=True,
              max_containers=150)
def verify(records: list) -> list:
    """records: [{"key", "content", "test"}] -> [{"key", "passed", "reason"}]"""
    import tempfile

    out = []
    for r in records:
        code = extract_code(r["content"])
        if code is None:
            out.append({"key": r["key"], "passed": False, "reason": "no code block"})
            continue
        with tempfile.TemporaryDirectory() as wd:
            t0 = time.time()
            ok, why = run_one(code, r["test"], wd)
        out.append({"key": r["key"], "passed": ok, "reason": why[:300], "secs": round(time.time() - t0, 2)})
    return out