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

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from dovla_cil.data.lerobot_export import (  # noqa: E402
    LeRobotExportConfig,
    export_lerobot_style_dataset,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Export a DoVLA-CIL dataset to a dependency-light LeRobot-style JSONL interchange "
            "format for external VLA baselines."
        )
    )
    parser.add_argument("--dataset", required=True, help="Source CIL dataset or collection dir.")
    parser.add_argument("--out", required=True, help="Output export directory.")
    parser.add_argument("--split", default="train", help="JSONL split filename stem.")
    parser.add_argument(
        "--selection",
        default="best",
        choices=["best", "expert", "success"],
        help="Which record to export from each CIL group.",
    )
    parser.add_argument("--max-groups", type=int, default=None)
    parser.add_argument(
        "--group-sampling",
        default="sequential",
        choices=["sequential", "task_balanced"],
        help="Ordering used before applying --max-groups.",
    )
    parser.add_argument("--seed", type=int, default=0, help="Seed for balanced group sampling.")
    parser.add_argument(
        "--no-images",
        action="store_true",
        help="Do not copy RGB frames even when observation_ref values are present.",
    )
    parser.add_argument("--image-format", default="jpg", choices=["jpg", "jpeg", "png"])
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    metadata = export_lerobot_style_dataset(
        LeRobotExportConfig(
            dataset_dir=args.dataset,
            out_dir=args.out,
            split=args.split,
            selection=args.selection,
            max_groups=args.max_groups,
            group_sampling=args.group_sampling,
            seed=args.seed,
            copy_images=not args.no_images,
            image_format=args.image_format,
        )
    )
    print(json.dumps(metadata, indent=2, sort_keys=True))
    return 0


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