#!/usr/bin/env python3
"""Execution-graded code eval (HumanEval pass@1) against an OpenAI-compatible endpoint.

The model is asked to complete a function; we extract its code block, assemble it with
the task's hidden test, and RUN it under a layered sandbox:
  - reliability_guard: neuters destructive syscalls (os.system, subprocess, rmtree, unlink, ...)
  - sandbox-exec '(allow default)(deny network*)': blocks all network (macOS Seatbelt)
  - fresh temp cwd, PYTHONDONTWRITEBYTECODE, hard timeout w/ SIGKILL

ALWAYS run --selftest first: it executes the CANONICAL solutions and must score ~100%.
If it doesn't, the executor is broken and no model number means anything.

Usage:
  python3 run_code.py --selftest
  python3 run_code.py --base http://LLAMA_SWAP_HOST:PORT --model <distilled-id> --tag code-distilled --key $API_KEY
  python3 run_code.py --base http://LLAMA_SWAP_HOST:PORT --model <stock-id>     --tag code-stock     --key $API_KEY
  python3 run_code.py --compare results_code/code-distilled.jsonl results_code/code-stock.jsonl
"""
import json, os, re, sys, time, argparse, tempfile, subprocess, urllib.request
from concurrent.futures import ThreadPoolExecutor

HERE = os.path.dirname(__file__)
DATA = os.path.join(HERE, "data_code")
RES = os.path.join(HERE, "results_code"); os.makedirs(RES, exist_ok=True)

SYS = ("You are an expert Python programmer. Complete the function described by the user. "
       "You may briefly reason about edge cases, complexity, and approach, but you MUST end "
       "your reply with the complete, runnable function inside a single ```python code block. "
       "Include the full function definition (signature and body) and any needed imports. "
       "Do not include example usage, explanation after the code, or tests.")

# Neutered inside the executed subprocess, before any candidate code runs.
GUARD = r"""
import os, shutil, subprocess, builtins, faulthandler
faulthandler.disable()
builtins.exit = None
builtins.quit = None
for _n in ("system","popen","remove","removedirs","rmdir","unlink","rename","renames",
           "replace","truncate","kill","killpg","fork","forkpty","chmod","chown","chroot"):
    if hasattr(os, _n): setattr(os, _n, None)
shutil.rmtree = None; shutil.move = None
subprocess.Popen = None; subprocess.run = None; subprocess.call = None
"""

def load_items():
    with open(os.path.join(DATA, "humaneval.jsonl")) as f:
        return [json.loads(l) for l in f if l.strip()]

# ---------- generation ----------
def call(base, model, key, prompt, max_tokens, temp, timeout, think=False):
    payload = {"model": model, "temperature": temp, "top_p": 1.0, "max_tokens": max_tokens,
               "messages": [{"role": "system", "content": SYS},
                            {"role": "user", "content": prompt}]}
    if think:
        payload["chat_template_kwargs"] = {"enable_thinking": True}
    req = urllib.request.Request(base.rstrip("/") + "/v1/chat/completions",
                                 data=json.dumps(payload).encode(),
                                 headers={"Content-Type": "application/json",
                                          **({"Authorization": f"Bearer {key}"} if key else {})})
    t0 = time.time()
    with urllib.request.urlopen(req, timeout=timeout) as r:
        d = json.load(r)
    dt = time.time() - t0
    ch = d["choices"][0]
    msg = ch["message"]
    text = msg.get("content") or msg.get("reasoning_content") or ""
    usage = d.get("usage", {})
    return text, usage.get("completion_tokens"), ch.get("finish_reason"), dt, d.get("model")

def extract_code(text):
    blocks = re.findall(r"```(?:python|py)?\s*\n(.*?)```", text, re.S)
    if blocks:
        return blocks[-1]
    m = re.search(r"(?:^|\n)(?:import |from |def |class )", text)
    return text[m.start():] if m else text

# ---------- execution ----------
def build_program(item, code):
    ep = item["entry_point"]
    has_def = re.search(r"\bdef\s+" + re.escape(ep) + r"\b", code)
    body = code if has_def else (item["prompt"] + "\n" + code)
    return body + "\n\n" + item["test"] + f"\n\ncheck({ep})\n"

def execute(program, timeout, use_sandbox=True):
    src = GUARD + "\n" + program
    with tempfile.TemporaryDirectory() as td:
        fp = os.path.join(td, "cand.py")
        with open(fp, "w") as f:
            f.write(src)
        py = [sys.executable, "-I", "-B", fp]
        cmd = (["sandbox-exec", "-p", "(version 1)(allow default)(deny network*)"] + py) if use_sandbox else py
        env = {**os.environ, "PYTHONDONTWRITEBYTECODE": "1", "OMP_NUM_THREADS": "1"}
        try:
            r = subprocess.run(cmd, cwd=td, env=env, timeout=timeout,
                               stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
            return (r.returncode == 0), ("" if r.returncode == 0 else r.stderr.decode("utf-8", "replace")[-240:])
        except subprocess.TimeoutExpired:
            return False, "timeout"
        except Exception as e:
            return False, f"exec-error: {e}"

# ---------- self-test: canonical solutions must pass ----------
def selftest(timeout, use_sandbox):
    items = load_items()
    print(f"[selftest] executing {len(items)} CANONICAL solutions "
          f"(sandbox={'on' if use_sandbox else 'off'})")
    ok = 0; fails = []
    for it in items:
        prog = build_program(it, it["canonical_solution"])
        passed, err = execute(prog, timeout, use_sandbox)
        ok += passed
        if not passed: fails.append((it["task_id"], err))
    print(f"[selftest] {ok}/{len(items)} canonical solutions pass")
    for tid, err in fails[:10]:
        print(f"   FAIL {tid}: {err}")
    if ok < len(items):
        print("\n!! executor not clean — fix before trusting model numbers.")
        if use_sandbox:
            print("   (retry with --no-sandbox to check if Seatbelt is the cause)")
    return ok == len(items)

# ---------- run a model ----------
def run(args):
    items = load_items()
    if args.limit: items = items[:args.limit]
    out = os.path.join(RES, f"{args.tag}.jsonl")
    print(f"[{args.tag}] {len(items)} items -> {args.model} @ {args.base}")

    def work(it):
        try:
            text, ctoks, finish, dt, served = call(args.base, args.model, args.key, it["prompt"],
                                                    args.max_tokens, args.temp, args.timeout, args.think)
            code = extract_code(text)
            passed, err = execute(build_program(it, code), args.exec_timeout, not args.no_sandbox)
            return {"task_id": it["task_id"], "entry_point": it["entry_point"],
                    "passed": bool(passed), "finish_reason": finish, "exec_err": err,
                    "completion_tokens": ctoks, "latency_s": round(dt, 2),
                    "served_model": served, "raw": text}
        except Exception as e:
            return {"task_id": it["task_id"], "entry_point": it["entry_point"],
                    "passed": False, "error": str(e)}

    rows = []
    with ThreadPoolExecutor(max_workers=args.concurrency) as ex:
        for i, r in enumerate(ex.map(work, items), 1):
            rows.append(r)
            flag = "ok " if r["passed"] else ("ERR" if r.get("error") else "x  ")
            trunc = " [TRUNC]" if r.get("finish_reason") == "length" else ""
            print(f"  [{i:>3}/{len(items)}] {flag} {r['task_id']:<16} {r.get('exec_err','')[:40]}{trunc}")
    with open(out, "w") as f:
        for r in rows: f.write(json.dumps(r) + "\n")

    served = sorted({r.get("served_model") for r in rows if r.get("served_model")})
    npass = sum(r["passed"] for r in rows)
    trunc = sum(1 for r in rows if r.get("finish_reason") == "length")
    avgtok = ([r["completion_tokens"] for r in rows if r.get("completion_tokens")] or [0])
    print(f"\nSERVED MODEL(S) this run: {served}")
    print(f"=== {args.tag} ===")
    print(f"pass@1: {npass}/{len(rows)} = {100*npass/len(rows):.1f}%")
    print(f"avg completion tokens: {sum(avgtok)//len(avgtok)}   truncated (finish=length): {trunc}")
    if trunc:
        print(f"  !! {trunc} items hit max_tokens — likely false zeros. Re-run with higher --max-tokens.")
    print(f"saved -> {out}")

def compare(a, b):
    def load(p): return {r["task_id"]: r for r in (json.loads(l) for l in open(p))}
    A, B = load(a), load(b)
    ta, tb = os.path.basename(a).split(".")[0], os.path.basename(b).split(".")[0]
    ids = [i for i in A if i in B]
    pa = sum(A[i]["passed"] for i in ids); pb = sum(B[i]["passed"] for i in ids)
    n = len(ids)
    print(f"{'metric':<16}{ta:>16}{tb:>16}")
    print(f"{'pass@1':<16}{100*pa/n:>15.1f}%{100*pb/n:>15.1f}%")
    print(f"{'solved':<16}{pa:>16}{pb:>16}  (N={n})")
    only_a = [i for i in ids if A[i]["passed"] and not B[i]["passed"]]
    only_b = [i for i in ids if B[i]["passed"] and not A[i]["passed"]]
    print(f"\nsolved only by {ta}: {len(only_a)}  -> {', '.join(only_a[:12])}")
    print(f"solved only by {tb}: {len(only_b)}  -> {', '.join(only_b[:12])}")
    sig = int(100 * (0.5 / (n ** 0.5)))
    print(f"\nN={n}. Rough ±{sig}pt 1-sigma per model; treat small gaps as directional.")

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--base"); ap.add_argument("--model"); ap.add_argument("--key")
    ap.add_argument("--tag", default="code-run")
    ap.add_argument("--max-tokens", type=int, default=6144, dest="max_tokens")
    ap.add_argument("--temp", type=float, default=0.0)
    ap.add_argument("--timeout", type=int, default=300, help="HTTP timeout per request")
    ap.add_argument("--exec-timeout", type=int, default=10, dest="exec_timeout", help="per-candidate run timeout")
    ap.add_argument("--think", action="store_true")
    ap.add_argument("--concurrency", type=int, default=4)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--no-sandbox", action="store_true", help="drop sandbox-exec (reliability_guard only)")
    ap.add_argument("--selftest", action="store_true")
    ap.add_argument("--compare", nargs=2)
    a = ap.parse_args()
    if a.selftest: sys.exit(0 if selftest(a.exec_timeout, not a.no_sandbox) else 1)
    elif a.compare: compare(*a.compare)
    else: run(a)