#!/usr/bin/env python3
"""Fetch HumanEval -> data_code/humaneval.jsonl.

Each row: task_id, prompt (signature+docstring), canonical_solution (body),
test (defines check()), entry_point (function name).

Usage:
  python3 prep_code.py            # all 164
  python3 prep_code.py 20         # first 20 (smoke test)
"""
import json, os, sys
from datasets import load_dataset

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

def main():
    n = int(sys.argv[1]) if len(sys.argv) > 1 else 0  # 0 = all
    ds = load_dataset("openai/openai_humaneval", split="test")
    rows = [{"task_id": r["task_id"], "prompt": r["prompt"],
             "canonical_solution": r["canonical_solution"],
             "test": r["test"], "entry_point": r["entry_point"]} for r in ds]
    if n: rows = rows[:n]
    p = os.path.join(OUT, "humaneval.jsonl")
    with open(p, "w") as f:
        for r in rows: f.write(json.dumps(r) + "\n")
    print(f"wrote {len(rows)} -> {p}")

if __name__ == "__main__":
    main()