#!/usr/bin/env python3
"""HumanEval pass@1 (greedy, completion-style prompt) for text-only qwen3_5 checkpoints / our bundles.

python3 scripts/humaneval.py --model /root/models/Qwen3.5-9B --out reports/humaneval_base.json
python3 scripts/humaneval.py --model exports/jev-judge-qwen35-9b-v0.8 --out reports/humaneval_v08.json
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import tempfile
import time

import torch

sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src"))

from jev_judge.model import load_text_causal_lm  # noqa: E402

STOPS = ["\nclass ", "\ndef ", "\n#", "\nif __name__", "\nprint(", "\nassert "]


def load_problems() -> list[dict]:
    from datasets import load_dataset

    ds = load_dataset("openai/openai_humaneval", split="test")
    return [dict(r) for r in ds]


def truncate(completion: str) -> str:
    cut = len(completion)
    for s in STOPS:
        i = completion.find(s)
        if i != -1:
            cut = min(cut, i)
    return completion[:cut]


def run_check(program: str, timeout: float = 15.0) -> tuple[bool, str]:
    with tempfile.TemporaryDirectory() as d:
        path = os.path.join(d, "prog.py")
        with open(path, "w") as f:
            f.write(program)
        try:
            r = subprocess.run([sys.executable, path], cwd=d, capture_output=True, text=True, timeout=timeout)
            return r.returncode == 0, (r.stderr[-300:] if r.returncode else "")
        except subprocess.TimeoutExpired:
            return False, "timeout"


@torch.no_grad()
def generate_all(model, tok, prompts: list[str], batch_size: int, max_new_tokens: int) -> list[str]:
    tok.padding_side = "left"
    outs: list[str] = []
    order = sorted(range(len(prompts)), key=lambda i: len(prompts[i]))
    result = [""] * len(prompts)
    for s in range(0, len(order), batch_size):
        idx = order[s : s + batch_size]
        batch = [prompts[i] for i in idx]
        enc = tok(batch, return_tensors="pt", padding=True, add_special_tokens=False).to("cuda")
        with torch.autocast("cuda", dtype=torch.bfloat16):
            gen = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=tok.pad_token_id)
        for j, i in enumerate(idx):
            result[i] = tok.decode(gen[j, enc["input_ids"].shape[1]:], skip_special_tokens=True)
        print(f"  generated {min(s + batch_size, len(order))}/{len(order)}", flush=True)
    return result


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--batch-size", type=int, default=16)
    ap.add_argument("--max-new-tokens", type=int, default=384)
    args = ap.parse_args()

    from transformers import AutoTokenizer

    problems = load_problems()
    tok = AutoTokenizer.from_pretrained(args.model)
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
    model, _ = load_text_causal_lm(args.model)
    t0 = time.time()
    raw = generate_all(model, tok, [p["prompt"] for p in problems], args.batch_size, args.max_new_tokens)
    gen_s = time.time() - t0
    results = []
    passed = 0
    for p, r in zip(problems, raw):
        comp = truncate(r)
        program = p["prompt"] + comp + "\n\n" + p["test"] + "\n" + f"check({p['entry_point']})\n"
        ok, err = run_check(program)
        passed += ok
        results.append({"task_id": p["task_id"], "passed": ok, "completion": comp, "error": err})
    score = passed / len(problems)
    summary = {"model": args.model, "n": len(problems), "passed": passed, "pass@1": score, "generation_s": gen_s,
               "protocol": "greedy, completion-style prompt, stop at " + repr(STOPS), "results": results}
    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
    with open(args.out, "w") as f:
        json.dump(summary, f, indent=1)
    print(f"HumanEval pass@1 = {score:.4f} ({passed}/{len(problems)}) | generation {gen_s:.0f}s | -> {args.out}")


if __name__ == "__main__":
    main()