#!/usr/bin/env python3
import argparse
import json
import os
import re
import time
from pathlib import Path

import pyarrow.parquet as pq


def safe_name(value: str, fallback: str) -> str:
    name = Path(value or fallback).name or fallback
    name = re.sub(r"[^A-Za-z0-9._-]+", "_", name)
    return name or fallback


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset-dir", required=True)
    parser.add_argument("--image-root", required=True)
    parser.add_argument("--json-out", required=True)
    parser.add_argument("--batch-size", type=int, default=512)
    parser.add_argument("--progress-every", type=int, default=5000)
    parser.add_argument("--max-rows", type=int, default=0)
    args = parser.parse_args()

    dataset_dir = Path(args.dataset_dir).resolve()
    image_root = Path(args.image_root).resolve()
    json_out = Path(args.json_out).resolve()
    tmp_out = json_out.with_suffix(json_out.suffix + ".tmp")

    parquet_files = sorted(dataset_dir.rglob("*.parquet"))
    if not parquet_files:
        raise FileNotFoundError(f"No parquet files found under {dataset_dir}")

    image_root.mkdir(parents=True, exist_ok=True)
    json_out.parent.mkdir(parents=True, exist_ok=True)

    total = 0
    started = time.time()
    with tmp_out.open("w", encoding="utf-8") as fout:
        fout.write("[\n")
        first = True
        for parquet_path in parquet_files:
            part = parquet_path.parent.name
            shard = parquet_path.stem
            shard_dir = image_root / part / shard
            shard_dir.mkdir(parents=True, exist_ok=True)

            pf = pq.ParquetFile(parquet_path)
            local_idx = 0
            for batch in pf.iter_batches(
                batch_size=args.batch_size,
                columns=["id", "conversations", "image_relpath", "image"],
            ):
                for row in batch.to_pylist():
                    if args.max_rows and total >= args.max_rows:
                        break
                    relpath = row.get("image_relpath") or ""
                    basename = safe_name(relpath, f"{row.get('id') or total}.png")
                    image_path = shard_dir / f"{local_idx:06d}_{basename}"
                    image_obj = row.get("image") or {}
                    image_bytes = image_obj.get("bytes")
                    if image_bytes is None:
                        raise ValueError(f"Missing image bytes in {parquet_path} row {local_idx}")
                    if not image_path.exists() or image_path.stat().st_size != len(image_bytes):
                        image_path.write_bytes(image_bytes)

                    record = {
                        "id": row.get("id"),
                        "image": str(image_path),
                        "conversations": row.get("conversations"),
                    }
                    if first:
                        first = False
                    else:
                        fout.write(",\n")
                    json.dump(record, fout, ensure_ascii=False)

                    total += 1
                    local_idx += 1
                    if total % args.progress_every == 0:
                        elapsed = max(time.time() - started, 1e-6)
                        rate = total / elapsed
                        print(f"materialized_rows={total} rate={rate:.1f}/s elapsed={elapsed:.1f}s", flush=True)
                if args.max_rows and total >= args.max_rows:
                    break
            if args.max_rows and total >= args.max_rows:
                break
        fout.write("\n]\n")

    os.replace(tmp_out, json_out)
    elapsed = max(time.time() - started, 1e-6)
    print(f"done rows={total} json={json_out} images={image_root} elapsed={elapsed:.1f}s", flush=True)


if __name__ == "__main__":
    main()