#!/usr/bin/env python3
"""Convert a pinned openai/gsm8k snapshot into verl v0.6 GRPO Parquet files."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import sys
from datetime import datetime, timezone
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq


DATASET = "openai/gsm8k"
REVISION = "740312add88f781978c0658806c59bc2815b9866"
INSTRUCTION = ' Let\'s think step by step and output the final answer after "####".'
RAW_FILES = {
    "train": {
        "relative_path": "main/train-00000-of-00001.parquet",
        "sha256": "ea82612ea9582142387730c793eb67d3b12849002bc0b7fa6f8efafa7351419d",
        "rows": 7473,
    },
    "test": {
        "relative_path": "main/test-00000-of-00001.parquet",
        "sha256": "ee7b8da9e381df27b9e3f7758a159ab2bdaa4dbaa910546cbbc47e0cb44e4f59",
        "rows": 1319,
    },
}


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def extract_solution(answer: str) -> str:
    match = re.search(r"#### (\-?[0-9\.\,]+)", answer)
    if match is None:
        raise ValueError(f"GSM8K answer does not contain a final solution: {answer!r}")
    return match.group(1).replace(",", "")


def convert_split(raw_root: Path, work_root: Path, split: str) -> dict[str, object]:
    spec = RAW_FILES[split]
    source = raw_root / str(spec["relative_path"])
    if not source.is_file():
        raise FileNotFoundError(source)
    observed_hash = sha256(source)
    if observed_hash != spec["sha256"]:
        raise RuntimeError(f"Raw {split} SHA-256 mismatch: {observed_hash}")

    raw_table = pq.read_table(source, columns=["question", "answer"])
    if raw_table.num_rows != spec["rows"]:
        raise RuntimeError(f"Unexpected {split} row count: {raw_table.num_rows}")

    records = []
    for index, raw in enumerate(raw_table.to_pylist()):
        question = raw["question"]
        answer = raw["answer"]
        if not isinstance(question, str) or not question:
            raise ValueError(f"Invalid question at {split}:{index}")
        if not isinstance(answer, str) or not answer:
            raise ValueError(f"Invalid answer at {split}:{index}")
        records.append(
            {
                "data_source": DATASET,
                "prompt": [{"role": "user", "content": question + INSTRUCTION}],
                "ability": "math",
                "reward_model": {"style": "rule", "ground_truth": extract_solution(answer)},
                "extra_info": {
                    "split": split,
                    "index": index,
                    "answer": answer,
                    "question": question,
                },
            }
        )

    output = work_root / f"{split}.parquet"
    table = pa.Table.from_pylist(records)
    pq.write_table(table, output, compression="snappy", use_dictionary=True)
    restored = pq.read_table(output)
    required_columns = {"data_source", "prompt", "ability", "reward_model", "extra_info"}
    if restored.num_rows != spec["rows"] or set(restored.column_names) != required_columns:
        raise RuntimeError(f"Prepared {split} Parquet failed schema or row-count validation")
    sample = restored.slice(0, 1).to_pylist()[0]
    if sample["data_source"] != DATASET or sample["reward_model"]["style"] != "rule":
        raise RuntimeError(f"Prepared {split} Parquet failed semantic validation")

    return {
        "file": output.name,
        "bytes": output.stat().st_size,
        "rows": restored.num_rows,
        "sha256": sha256(output),
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--raw-root", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    args = parser.parse_args()

    raw_root = args.raw_root.resolve()
    output_root = args.output_root.resolve()
    source = json.loads((raw_root / "SOURCE.json").read_text(encoding="utf-8"))
    if source.get("dataset") != DATASET or source.get("revision") != REVISION:
        raise RuntimeError("Raw dataset identity does not match the pinned GSM8K snapshot")
    if output_root.exists():
        raise FileExistsError(f"Refusing to overwrite prepared dataset: {output_root}")

    work_root = output_root.with_name(f"{output_root.name}.preparing-{os.getpid()}")
    if work_root.exists():
        raise FileExistsError(work_root)
    work_root.mkdir(parents=True)
    try:
        outputs = {split: convert_split(raw_root, work_root, split) for split in ("train", "test")}
        manifest = {
            "status": "prepared",
            "prepared_at": datetime.now(timezone.utc).isoformat(),
            "dataset": DATASET,
            "revision": REVISION,
            "configuration": "main",
            "conversion_authority": "verl v0.6.0 examples/data_preprocess/gsm8k.py",
            "conversion_script_sha256": sha256(Path(__file__).resolve()),
            "python": sys.version.split()[0],
            "pyarrow": pa.__version__,
            "raw_files": RAW_FILES,
            "outputs": outputs,
        }
        (work_root / "dataset-manifest.json").write_text(
            json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
        )
        checksum_lines = [f"{outputs[split]['sha256']}  {outputs[split]['file']}" for split in ("train", "test")]
        (work_root / "SHA256SUMS").write_text("\n".join(checksum_lines) + "\n", encoding="utf-8")
        os.replace(work_root, output_root)
    except BaseException:
        shutil.rmtree(work_root, ignore_errors=True)
        raise

    print(json.dumps(manifest, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())