import json
from pathlib import Path

def adapt_mbpp_task(entry: dict) -> dict:
    """Convert MBPP entry to CodeForge training format."""
    task_id = entry["task_id"]
    brief = entry["text"]

    # Starter file: empty main.py
    initial_files = {"main.py": "# Write your solution here\n"}

    # Hidden tests: wrap MBPP asserts into pytest functions
    test_lines = []
    setup = entry.get("test_setup_code", "")
    if setup:
        test_lines.append(setup)
    test_lines.append("from main import *\n")
    for i, assert_str in enumerate(entry.get("test_list", [])):
        test_lines.append(f"def test_{i}():\n    {assert_str}\n")

    hidden_tests = {"test_hidden.py": "\n".join(test_lines)}

    return {
        "task_id": str(task_id),
        "brief": brief,
        "initial_files": initial_files,
        "hidden_tests": hidden_tests,
        "tools": ("ruff", "imports", "mypy", "pytest"),
    }

def load_mbpp_dataset(dataset_path: str = "dataset/mbpp.jsonl", split: str = "train"):
    """Load MBPP tasks and map them to CodeForge format suitable for TRL."""
    import datasets
    
    tasks = []
    path = Path(dataset_path)
    if not path.exists():
        raise FileNotFoundError(f"MBPP dataset not found at {path}")
        
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            if not line.strip(): continue
            tasks.append(json.loads(line))
            
    # Split: 1-874 Train, 875-974 Eval
    if split == "train":
        tasks = tasks[:874]
    else:
        tasks = tasks[874:]
        
    formatted_tasks = [adapt_mbpp_task(t) for t in tasks]
    
    # Create huggingface dataset
    def gen():
        for t in formatted_tasks:
            yield {
                "prompt": [
                    {"role": "system", "content": "You are a Python developer. Solve this problem. Return ONLY the Python code for main.py, no markdown fences."},
                    {"role": "user", "content": f"## Problem\n{t['brief']}\n\n## Rules\n- Write your solution in a single `main.py`\n- All functions must have type hints\n- Code must pass ruff, mypy, and pytest"}
                ],
                "task_meta": t
            }
            
    return datasets.Dataset.from_generator(gen)