"""Download a LeRobot dataset from HuggingFace Hub and prepare it for the dataloader.

This script:
1. Downloads a LeRobot dataset from HuggingFace Hub.
2. Loads it into Rerun via the built-in LeRobot importer (`log_file_from_path`).
3. Splits the resulting archive into one RRD per episode.
4. Registers the per-episode RRDs to a catalog server instance.

"""

from __future__ import annotations

import argparse
import re
from pathlib import Path

from huggingface_hub import snapshot_download

import rerun as rr

DEFAULT_REPO_ID = "rerun/so101-pick-and-place"
DEFAULT_OUTPUT_DIR = Path(__file__).resolve().parent / "data"
APPLICATION_ID = "lerobot"

_EPISODE_RE = re.compile(r"^(episode_)(\d+)$")


def _zero_pad_episode_id(rec_id: str, width: int = 5) -> str:
    """Turn `episode_1` into `episode_00001` so segments sort lexicographically."""
    m = _EPISODE_RE.match(rec_id)
    if m:
        return f"{m.group(1)}{int(m.group(2)):0{width}d}"
    return rec_id


def download_dataset(repo_id: str, dest: Path) -> Path:
    """Download a LeRobot dataset from HuggingFace Hub into *dest* and return its path."""
    print(f"Downloading {repo_id} to {dest} …")
    local_dir = snapshot_download(repo_id=repo_id, repo_type="dataset", local_dir=dest)
    return Path(local_dir)


def lerobot_to_combined_rrd(dataset_dir: Path, combined_rrd: Path) -> None:
    """Use Rerun's built-in LeRobot importer to turn the dataset into a single RRD."""
    print(f"Converting {dataset_dir} -> {combined_rrd}")
    with rr.RecordingStream(APPLICATION_ID) as rec:
        rec.save(str(combined_rrd))
        rec.log_file_from_path(str(dataset_dir))


def split_into_episode_rrds(combined_rrd: Path, rrd_dir: Path) -> list[Path]:
    """Split a combined RRD archive into one RRD per episode.

    Returns the paths of the written per-episode RRDs.
    """
    rrd_dir.mkdir(parents=True, exist_ok=True)

    reader = rr.experimental.RrdReader(str(combined_rrd))
    recordings = reader.recordings()
    print(f"Archive contains {len(recordings)} recordings")

    episode_paths: list[Path] = []
    for entry in recordings:
        store = reader.store(store=entry)
        # Skip metadata-only recordings (e.g. the "root" recording that only carries properties).
        if not store.schema().entity_paths():
            continue

        episode_id = _zero_pad_episode_id(entry.recording_id)
        rrd_path = rrd_dir / f"{episode_id}.rrd"

        with rr.RecordingStream(APPLICATION_ID, recording_id=episode_id, send_properties=False) as rec:
            rec.save(str(rrd_path))
            rec.send_chunks(store)
        episode_paths.append(rrd_path)
        print(f"  wrote {rrd_path} ({rrd_path.stat().st_size / (1024 * 1024):.1f} MB)")

    return episode_paths


def register_to_catalog(
    rrd_paths: list[Path],
    *,
    catalog_url: str,
    dataset_name: str,
) -> None:
    """Register per-episode RRDs to a catalog server instance.

    Uses absolute file:// URIs so the catalog can read the RRDs directly from the local filesystem.
    """
    print(f"\nRegistering {len(rrd_paths)} episodes to {catalog_url} as dataset '{dataset_name}' …")
    client = rr.catalog.CatalogClient(catalog_url)
    dataset = client.create_dataset(dataset_name, exist_ok=True)

    uris = [f"file://{p.resolve()}" for p in rrd_paths]
    on_duplicate = rr.catalog.OnDuplicateSegmentLayer(rr.catalog.OnDuplicateSegmentLayer.REPLACE)
    dataset.register(uris, on_duplicate=on_duplicate).wait()
    print("  registration done")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument(
        "--repo-id",
        default=DEFAULT_REPO_ID,
        help=f"HuggingFace dataset repo id (default: {DEFAULT_REPO_ID}).",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=DEFAULT_OUTPUT_DIR,
        help=f"Directory to store downloaded dataset and output RRDs (default: {DEFAULT_OUTPUT_DIR}).",
    )
    parser.add_argument(
        "--catalog-url",
        default="rerun+http://127.0.0.1:51234",
        help="Rerun catalog URL to register episodes with. Pass an empty string to skip registration.",
    )
    parser.add_argument(
        "--dataset-name",
        default=None,
        help="Name of the dataset to create/use in the catalog (default: derived from --repo-id).",
    )
    parser.add_argument(
        "--keep-combined",
        action="store_true",
        help="Keep the intermediate combined RRD after splitting (useful for debugging).",
    )
    args = parser.parse_args()

    output_dir = args.output_dir
    repo_slug = args.repo_id.replace("/", "_")
    dataset_dir = output_dir / "lerobot" / repo_slug
    rrd_dir = output_dir / "rrds" / repo_slug
    combined_rrd = output_dir / f"{repo_slug}_combined.rrd"

    download_dataset(args.repo_id, dataset_dir)
    lerobot_to_combined_rrd(dataset_dir, combined_rrd)
    episode_paths = split_into_episode_rrds(combined_rrd, rrd_dir)

    if not args.keep_combined:
        combined_rrd.unlink()

    print(f"\nWrote {len(episode_paths)} per-episode RRDs to {rrd_dir}")

    if args.catalog_url:
        dataset_name = args.dataset_name or repo_slug
        register_to_catalog(episode_paths, catalog_url=args.catalog_url, dataset_name=dataset_name)


if __name__ == "__main__":
    main()