"""Download livecodebench/code_generation_lite (v6, 2501-2505)
from HuggingFace and convert to the agentic-rl training format.

This script keeps only contest problems whose contest_date falls in
[2025-01-01, 2025-06-01), matching the demo validation subset.

Target format (one JSON object per line):
{
  "messages": [{"content": "<problem>", "role": "user"}],
  "rollout_extra": {
      "public_test_cases": "...",
      "private_test_cases": "...",
      "starter_code": "...",
      "platform": "...",
      "question_id": "...",
      "difficulty": "...",
      "contest_date": "...",
      "dataset": "livecodebench_v6"
  }
}
"""

import json
from pathlib import Path

from datasets import load_dataset

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

VERSION_TAG = "release_v6"
CONFIG_NAME = "release_latest"
START_DATE = "2025-01-01"
END_DATE = "2025-06-01"


def is_target_contest_date(example: dict) -> bool:
    contest_date = str(example.get("contest_date", "")).strip()[:10]
    return START_DATE <= contest_date < END_DATE


def convert_example(example: dict) -> dict:
    return {
        "messages": [
            {"content": example["question_content"], "role": "user"}
        ],
        "rollout_extra": {
            "question_title": example["question_title"],
            "public_test_cases": example["public_test_cases"],
            "private_test_cases": example["private_test_cases"],
            "starter_code": example["starter_code"],
            "platform": example["platform"],
            "question_id": example["question_id"],
            "contest_id": example["contest_id"],
            "contest_date": example["contest_date"],
            "difficulty": example["difficulty"],
            "metadata": example["metadata"],
            "dataset": "livecodebench_v6",
        },
    }


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

    print(f"Loading livecodebench/code_generation_lite ({CONFIG_NAME}, {VERSION_TAG}) ...")
    ds = load_dataset(
        "livecodebench/code_generation_lite",
        CONFIG_NAME,
        version_tag=VERSION_TAG,
        trust_remote_code=True,
    )

    raw_examples = list(ds["test"])
    filtered_examples = [ex for ex in raw_examples if is_target_contest_date(ex)]
    if not filtered_examples:
        raise ValueError(
            f"No LiveCodeBench examples found for {START_DATE} <= contest_date < {END_DATE}"
        )

    all_examples = [convert_example(ex) for ex in filtered_examples]

    out_path = OUTPUT_DIR / "livecodebench_v6_test.jsonl"
    with open(out_path, "w", encoding="utf-8") as f:
        for item in all_examples:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

    difficulty_counts = {}
    platform_counts = {}
    date_counts = {}
    for item in all_examples:
        d = item["rollout_extra"]["difficulty"]
        p = item["rollout_extra"]["platform"]
        month = str(item["rollout_extra"]["contest_date"])[:7]
        difficulty_counts[d] = difficulty_counts.get(d, 0) + 1
        platform_counts[p] = platform_counts.get(p, 0) + 1
        date_counts[month] = date_counts.get(month, 0) + 1

    print(f"\nDone!")
    print(f"  Raw test split: {len(raw_examples)} examples")
    print(f"  Date filter: {START_DATE} <= contest_date < {END_DATE}")
    print(f"  Total: {len(all_examples)} examples -> {out_path}")
    print(f"  Contest month distribution: {dict(sorted(date_counts.items()))}")
    print(f"  Difficulty distribution: {difficulty_counts}")
    print(f"  Platform distribution: {platform_counts}")


if __name__ == "__main__":
    main()
