#!/usr/bin/env python3
from __future__ import annotations

import argparse
import glob
import json
from pathlib import Path

import pyarrow.parquet as pq


def iter_paths(patterns: list[str]) -> list[Path]:
    paths: list[Path] = []
    for pattern in patterns:
        path = Path(pattern)
        if path.is_dir():
            paths.extend(path.glob("*.parquet"))
            continue
        matches = glob.glob(pattern)
        if matches:
            paths.extend(Path(match) for match in matches)
            continue
        if path.exists():
            paths.append(path)
            continue
        raise FileNotFoundError(pattern)
    return sorted(set(paths), key=str)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Export pretrain parquet text rows to Megatron JSONL.")
    parser.add_argument("--input", action="append", required=True, help="Parquet directory/file/glob.")
    parser.add_argument("--output", required=True, help="Output JSONL path.")
    parser.add_argument("--text-field", default="text")
    parser.add_argument("--batch-size", type=int, default=8192)
    parser.add_argument("--max-docs", type=int, default=0, help="0 means no limit.")
    parser.add_argument("--progress-every", type=int, default=25)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    paths = iter_paths(args.input)
    if not paths:
        raise SystemExit("no parquet inputs resolved")

    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    docs = 0
    with output.open("w", encoding="utf-8") as handle:
        for file_idx, path in enumerate(paths, start=1):
            parquet_file = pq.ParquetFile(path)
            if args.text_field not in parquet_file.schema_arrow.names:
                continue
            for batch in parquet_file.iter_batches(
                batch_size=args.batch_size,
                columns=[args.text_field],
                use_threads=True,
            ):
                for row in batch.to_pylist():
                    text = row.get(args.text_field)
                    if not isinstance(text, str) or not text.strip():
                        continue
                    handle.write(json.dumps({"text": text}, ensure_ascii=False) + "\n")
                    docs += 1
                    if args.max_docs and docs >= args.max_docs:
                        print(f"export_done docs={docs} files_seen={file_idx}/{len(paths)} output={output}")
                        return
            if file_idx % args.progress_every == 0:
                print(f"export_progress files={file_idx}/{len(paths)} docs={docs}")
    print(f"export_done docs={docs} files_seen={len(paths)} output={output}")


if __name__ == "__main__":
    main()
