#!/usr/bin/env python3
"""Convert parquet datasets in ``data/`` to JSONL format.

Each parquet row is expected to follow the schema produced by the calc dataset
pipeline::

    prompt        struct<content: string, role: string>
    reward_model  struct<ground_truth: string, style: string>
    data_source   string
    agent_name    string
    extra_info    struct<expression: string, index: int64, split: string>

The emitted JSONL records mirror the layout used by the existing ``rl-*.jsonl``
files in ``data/``::

    {"messages": [{"content": "...", "role": "user"}],
     "rollout_extra": {"solution": "..."}}

Usage::

    python scripts/convert_parquet_to_jsonl.py
    python scripts/convert_parquet_to_jsonl.py path/to/in.parquet path/to/out.jsonl
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Iterable

try:
    import pyarrow.parquet as pq
except ImportError as exc:  # pragma: no cover - runtime guard
    raise SystemExit(
        "pyarrow is required. Install it with `pip install pyarrow`."
    ) from exc


REPO_ROOT = Path(__file__).resolve().parent.parent
DATA_DIR = REPO_ROOT / "data"

DEFAULT_PAIRS: list[tuple[Path, Path]] = [
    (DATA_DIR / "train_convert.parquet", DATA_DIR / "calc_train.jsonl"),
    (DATA_DIR / "test_convert.parquet", DATA_DIR / "calc_validation.jsonl"),
]


def _row_to_record(row: dict) -> dict:
    prompt = row.get("prompt")
    reward_model = row.get("reward_model") or {}

    messages = prompt if isinstance(prompt, list) else [prompt]
    solution = reward_model.get("ground_truth", "")

    return {
        "messages": messages,
        "rollout_extra": {"solution": solution},
    }


def convert_file(src: Path, dst: Path) -> int:
    """Convert a single parquet file to JSONL. Returns the number of rows written."""
    if not src.is_file():
        raise FileNotFoundError(f"Parquet file not found: {src}")

    table = pq.read_table(src)
    dst.parent.mkdir(parents=True, exist_ok=True)

    written = 0
    with dst.open("w", encoding="utf-8") as fout:
        for row in table.to_pylist():
            record = _row_to_record(row)
            fout.write(json.dumps(record, ensure_ascii=False))
            fout.write("\n")
            written += 1
    return written


def _iter_pairs(args: argparse.Namespace) -> Iterable[tuple[Path, Path]]:
    if args.input is None and args.output is None:
        yield from DEFAULT_PAIRS
        return

    if args.input is None or args.output is None:
        raise SystemExit("Both INPUT and OUTPUT must be provided together.")

    yield Path(args.input), Path(args.output)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", nargs="?", help="Input .parquet file")
    parser.add_argument("output", nargs="?", help="Output .jsonl file")
    args = parser.parse_args()

    total = 0
    for src, dst in _iter_pairs(args):
        count = convert_file(src, dst)
        total += count
        print(f"[ok] {src} -> {dst} ({count} rows)")

    print(f"[done] wrote {total} rows in total")
    return 0


if __name__ == "__main__":
    sys.exit(main())
