"""MBPP sanitized load/filter for small-model coding RL."""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from deepagents_rl.rewards.unit_test import run_hidden_tests


def _row_to_task(ex: dict[str, Any], split: str) -> dict[str, Any] | None:
    """Normalize an MBPP sanitized example."""
    task_id = ex.get("task_id") or ex.get("id")
    text = ex.get("text") or ex.get("prompt") or ex.get("description") or ""
    code = ex.get("code") or ex.get("reference_code") or ""
    test_list = list(ex.get("test_list") or ex.get("test_cases") or [])
    if not text or not test_list:
        return None
    # Leakage-aware: first assert may appear in prompt as visible example
    visible = [test_list[0]] if test_list else []
    hidden = list(test_list)
    prompt = (
        f"Write a Python solution for the following problem.\n\n"
        f"{text.strip()}\n\n"
        f"Visible example test (not the full grader):\n{visible[0]}\n\n"
        f"Return ONLY a ```python``` code block with the full solution."
    )
    return {
        "task_id": str(task_id),
        "split": split,
        "text": text.strip(),
        "reference_code": code,
        "visible_tests": visible,
        "hidden_tests": hidden,
        "prompt": prompt,
        "test_list": test_list,
    }


def filter_solvable(tasks: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    """Keep tasks whose reference solution passes hidden tests in our sandbox."""
    keep: list[dict[str, Any]] = []
    drop: list[dict[str, Any]] = []
    for t in tasks:
        ref = t.get("reference_code") or ""
        if not ref.strip():
            drop.append({**t, "drop_reason": "no_reference"})
            continue
        grade = run_hidden_tests(ref, t["hidden_tests"], timeout_s=5.0)
        if grade["pass_frac"] >= 1.0:
            keep.append(t)
        else:
            drop.append({**t, "drop_reason": "ref_fail", "grade": grade})
    return keep, drop


def load_mbpp_sanitized() -> dict[str, list[dict[str, Any]]]:
    """Load google-research-datasets/mbpp sanitized config from HF."""
    from datasets import load_dataset

    # Prefer sanitized; fall back to full mbpp
    try:
        ds = load_dataset("google-research-datasets/mbpp", "sanitized")
    except Exception:
        ds = load_dataset("mbpp", "sanitized")

    out: dict[str, list[dict[str, Any]]] = {}
    for split in ds:
        tasks = []
        for ex in ds[split]:
            t = _row_to_task(dict(ex), split=split)
            if t:
                tasks.append(t)
        out[split] = tasks
    return out


def build_train_eval(
    *,
    filter_refs: bool = True,
) -> dict[str, list[dict[str, Any]]]:
    """
    Train = sanitized train + validation + prompt (none of test).
    Eval = sanitized test (held-out).
    """
    raw = load_mbpp_sanitized()
    train_parts = []
    for key in ("train", "validation", "prompt"):
        train_parts.extend(raw.get(key, []))
    eval_parts = list(raw.get("test", []))

    drop_log: list[dict[str, Any]] = []
    if filter_refs:
        train_parts, d1 = filter_solvable(train_parts)
        eval_parts, d2 = filter_solvable(eval_parts)
        drop_log = d1 + d2

    # Disjoint task ids
    train_ids = {t["task_id"] for t in train_parts}
    eval_parts = [t for t in eval_parts if t["task_id"] not in train_ids]

    return {
        "train": train_parts,
        "eval": eval_parts,
        "dropped": drop_log,
    }


def save_jsonl(tasks: list[dict[str, Any]], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for t in tasks:
            # Don't need full drop grades in published rows
            row = {k: v for k, v in t.items() if k not in ("drop_reason", "grade")}
            f.write(json.dumps(row) + "\n")


def to_hf_dataset(splits: dict[str, list[dict[str, Any]]]):
    from datasets import Dataset, DatasetDict

    def _clean(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
        cleaned = []
        for t in rows:
            cleaned.append(
                {
                    "task_id": t["task_id"],
                    "split": t["split"],
                    "text": t["text"],
                    "prompt": t["prompt"],
                    "reference_code": t.get("reference_code", ""),
                    "visible_tests": t["visible_tests"],
                    "hidden_tests": t["hidden_tests"],
                }
            )
        return cleaned

    return DatasetDict(
        {
            "train": Dataset.from_list(_clean(splits["train"])),
            "eval": Dataset.from_list(_clean(splits["eval"])),
        }
    )