#!/usr/bin/env python3
"""HumanEval pass@1 against llama-server (OpenAI-compatible chat API).

Greedy decoding (temperature 0), max 2048 new tokens. Start the server with thinking turned off:
    ./serve_b200.sh --reasoning-budget 0

The model reply is expected to contain the full function in a ```python block; the code is run against the official
HumanEval tests in a subprocess (30 s, 8 GiB address-space limit).
WARNING: this executes model-generated code on your machine. Run it in a sandbox or a throw-away container.

Usage:
    pip install aiohttp datasets
    python humaneval.py --url http://127.0.0.1:8080 --concurrency 1 --out humaneval_results.json
"""
import argparse, asyncio, json, os, re, resource, subprocess, sys, tempfile, time
from concurrent.futures import ThreadPoolExecutor

import aiohttp

PROMPT = ("Complete the following Python function. Reply with the complete function (including any imports it needs) "
          "in a single ```python code block.\n\n```python\n{prompt}```")


async def chat(session, url, content, max_tokens, sem):
    body = {"messages": [{"role": "user", "content": content}], "temperature": 0, "max_tokens": max_tokens}
    async with sem:
        t0 = time.time()
        async with session.post(f"{url}/v1/chat/completions", json=body, timeout=aiohttp.ClientTimeout(total=3600)) as r:
            j = await r.json()
    ch = j["choices"][0]
    m = ch["message"]
    return {"content": m.get("content") or "", "reasoning": m.get("reasoning_content") or "",
            "finish": ch.get("finish_reason"), "ntok": j.get("usage", {}).get("completion_tokens", 0), "sec": time.time() - t0}


def extract_code(resp, entry):
    blocks = re.findall(r"```(?:python|py|Python)?[ \t]*\n(.*?)```", resp, re.S)
    if not blocks:
        return resp if "def " in resp else ""
    for b in blocks:
        if re.search(rf"def\s+{re.escape(entry)}\s*\(", b):
            return b
    return max(blocks, key=len)


def limits():
    resource.setrlimit(resource.RLIMIT_AS, (8 << 30, 8 << 30))
    resource.setrlimit(resource.RLIMIT_CPU, (30, 30))


def run_program(prog, timeout=30):
    with tempfile.TemporaryDirectory() as d:
        p = os.path.join(d, "t.py")
        open(p, "w").write(prog)
        try:
            r = subprocess.run([sys.executable, p], cwd=d, capture_output=True, text=True, timeout=timeout,
                               preexec_fn=limits, env={"PATH": os.environ.get("PATH", ""), "HOME": d})
            return r.returncode == 0, (r.stderr or "")[-300:]
        except subprocess.TimeoutExpired:
            return False, "timeout"


async def main(a):
    from datasets import load_dataset

    items = list(load_dataset("openai/openai_humaneval", split="test"))
    items = items[: a.limit] if a.limit else items
    sem = asyncio.Semaphore(a.concurrency)
    done = 0

    async def one(sess, it, pool):
        nonlocal done
        r = await chat(sess, a.url, PROMPT.format(prompt=it["prompt"]), a.max_tokens, sem)
        code = extract_code(r["content"], it["entry_point"])
        if re.search(rf"def\s+{re.escape(it['entry_point'])}\s*\(", code):
            prog = it["prompt"] + "\n    pass\n\n" + code + "\n\n" + it["test"] + f"\n\ncheck({it['entry_point']})\n"
        else:
            prog = it["prompt"] + code + "\n\n" + it["test"] + f"\n\ncheck({it['entry_point']})\n"
        ok, err = await asyncio.get_running_loop().run_in_executor(pool, run_program, prog)
        done += 1
        if done % 10 == 0:
            print(f"  {done}/{len(items)}", flush=True)
        return {"task_id": it["task_id"], "passed": ok, "err": err, **r, "code": code}

    with ThreadPoolExecutor(8) as pool:
        async with aiohttp.ClientSession() as sess:
            t0 = time.time()
            res = await asyncio.gather(*[one(sess, it, pool) for it in items])
    n = len(res); p = sum(r["passed"] for r in res)
    print(f"HumanEval pass@1 = {p}/{n} = {100 * p / n:.1f}%  (hit max tokens {sum(r['finish'] == 'length' for r in res)}, "
          f"no code {sum(not r['code'] for r in res)}, avg {sum(r['ntok'] for r in res) / n:.0f} completion tokens, {time.time() - t0:.0f}s)")
    json.dump(res, open(a.out, "w"), indent=1)
    for r in res:
        if not r["passed"]:
            last = r["err"].strip().splitlines()[-1][:110] if r["err"].strip() else r["finish"]
            print("  FAIL", r["task_id"], "|", last)


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--url", default="http://127.0.0.1:8080", help="llama-server base URL")
    ap.add_argument("--out", default="humaneval_results.json")
    ap.add_argument("--max-tokens", type=int, default=2048)
    ap.add_argument("--concurrency", type=int, default=1, help="parallel requests (1 reproduces the reported number)")
    ap.add_argument("--limit", type=int, default=0, help="only run the first N problems")
    asyncio.run(main(ap.parse_args()))