"""Download EleutherAI/hendrycks_math from HuggingFace and convert to
the agentic-rl training format.

Target format (one JSON object per line):
{
  "messages": [{"content": "<problem> ...", "role": "user"}],
  "rollout_extra": {"solution": "<boxed_answer>"},
  "ground_truth": "<boxed_answer>"
}
"""

import json
import re
import sys
from pathlib import Path

from datasets import load_dataset

OUTPUT_DIR = Path(__file__).resolve().parent.parent / "data"
DATASET_REVISION = "21a5633873b6a120296cce3e2df9d5550074f4a3"

SUBSETS = [
    "algebra",
    "counting_and_probability",
    "geometry",
    "intermediate_algebra",
    "number_theory",
    "prealgebra",
    "precalculus",
]

def extract_boxed_answer(solution: str) -> str:
    """Extract the content inside the last \\boxed{...} in the solution.

    Handles nested braces correctly, e.g. \\boxed{\\frac{1}{2}}.
    """
    idx = solution.rfind("\\boxed{")
    if idx == -1:
        idx = solution.rfind("\\boxed ")
        if idx == -1:
            return "0"
        start = idx + len("\\boxed ")
        match = re.match(r"\S+", solution[start:])
        return match.group(0) if match else "0"

    start = idx + len("\\boxed{")
    depth = 1
    i = start
    while i < len(solution) and depth > 0:
        if solution[i] == "{":
            depth += 1
        elif solution[i] == "}":
            depth -= 1
        i += 1
    answer = solution[start : i - 1]
    return answer if answer else "0"


def convert_example(example: dict) -> dict:
    problem = example["problem"]
    answer = extract_boxed_answer(example["solution"])
    return {
        "messages": [
            {"content": problem, "role": "user"}
        ],
        "rollout_extra": {
            "solution": answer,
            "dataset": "hendrycks_math",
            "level": example.get("level", ""),
            "type": example.get("type", ""),
        },
        "ground_truth": answer,
    }


def is_level_45(item: dict) -> bool:
    level = str(item["rollout_extra"].get("level", "")).strip().lower()
    return level in {"level 4", "level 5", "4", "5"}


def write_jsonl(records: list[dict], path: Path) -> None:
    with open(path, "w", encoding="utf-8") as f:
        for item in records:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")


def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    all_train = []
    all_test = []

    for subset in SUBSETS:
        print(f"Loading subset: {subset} ...")
        ds = load_dataset("EleutherAI/hendrycks_math", subset, revision=DATASET_REVISION)

        for example in ds["train"]:
            all_train.append(convert_example(example))
        for example in ds["test"]:
            all_test.append(convert_example(example))

    train_path = OUTPUT_DIR / "hendrycks_math_train.jsonl"
    test_path = OUTPUT_DIR / "hendrycks_math_test.jsonl"
    train_lv45_path = OUTPUT_DIR / "hendrycks_math_train_lv45.jsonl"
    test_lv45_path = OUTPUT_DIR / "hendrycks_math_test_lv45.jsonl"

    train_lv45 = [item for item in all_train if is_level_45(item)]
    test_lv45 = [item for item in all_test if is_level_45(item)]

    write_jsonl(all_train, train_path)
    write_jsonl(all_test, test_path)
    write_jsonl(train_lv45, train_lv45_path)
    write_jsonl(test_lv45, test_lv45_path)

    empty_answers = sum(
        1 for item in all_train + all_test if not item["rollout_extra"]["solution"]
    )

    print(f"\nDone!")
    print(f"  Full train:      {len(all_train)} examples -> {train_path}")
    print(f"  Full test:       {len(all_test)} examples -> {test_path}")
    print(f"  Level 4-5 train: {len(train_lv45)} examples -> {train_lv45_path}")
    print(f"  Level 4-5 test:  {len(test_lv45)} examples -> {test_lv45_path}")
    if empty_answers:
        print(f"  Warning: {empty_answers} examples had empty boxed answers")


if __name__ == "__main__":
    main()
