#!/usr/bin/env python3
"""HumanEval pass@1 through vLLM (AR path of a dual-head bundle). Same prompts / truncation / checker as humaneval.py.

python3 scripts/vllm_humaneval.py --model exports/jev-judge-qwen35-9b-v0.8 --out reports/humaneval_vllm_9b.json
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time

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

from humaneval import load_problems, run_check, truncate  # noqa: E402


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

    from vllm import LLM, SamplingParams

    problems = load_problems()
    t0 = time.time()
    llm = LLM(model=args.model, dtype="bfloat16", gpu_memory_utilization=args.gpu_mem, max_model_len=4096,
              enable_prefix_caching=False)
    load_s = time.time() - t0
    sp = SamplingParams(temperature=0.0, max_tokens=args.max_new_tokens)
    t0 = time.time()
    outs = llm.generate([p["prompt"] for p in problems], sp, use_tqdm=False)
    gen_s = time.time() - t0
    n_tok = sum(len(o.outputs[0].token_ids) for o in outs)
    results, passed = [], 0
    for p, o in zip(problems, outs):
        comp = truncate(o.outputs[0].text)
        ok, err = run_check(p["prompt"] + comp + "\n\n" + p["test"] + "\n" + f"check({p['entry_point']})\n")
        passed += ok
        results.append({"task_id": p["task_id"], "passed": ok, "completion": comp, "error": err})
    summary = {"model": args.model, "engine": "vllm", "n": len(problems), "passed": passed, "pass@1": passed / len(problems),
               "generation_s": gen_s, "generated_tokens": n_tok, "gen_tok_per_s": n_tok / gen_s, "load_s": load_s, "results": results}
    with open(args.out, "w") as f:
        json.dump(summary, f, indent=1)
    print(f"vLLM HumanEval pass@1 = {passed/len(problems):.4f} ({passed}/{len(problems)}) | generation {gen_s:.1f}s "
          f"({n_tok/gen_s:.0f} generated tok/s) | load {load_s:.0f}s -> {args.out}")


if __name__ == "__main__":
    main()