#!/usr/bin/env python
"""Download Stage-2 dataset subsets under a hard disk budget.

Usage:
    python scripts/download_subsets.py --root ~/tinyvla_data [--dry-run]

Downloads metadata first, measures actual on-disk size of a small episode
sample, then extrapolates before committing to the full subset download.
"""

from __future__ import annotations

import argparse
import shutil
from pathlib import Path

DISK_BUDGET_GB = 400

# (repo_id, num_episodes or None for all)
SUBSETS = [
    ("HuggingFaceVLA/community_dataset_v1", None),
    ("nvidia/BridgeData2_LeRobot_v3", 10_000),
    ("IPEC-COMMUNITY/fractal20220817_data_lerobot", 8_000),
]


def free_gb(path: Path) -> float:
    return shutil.disk_usage(path).free / 1e9


def used_gb(path: Path) -> float:
    return sum(f.stat().st_size for f in path.rglob("*") if f.is_file()) / 1e9


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, default=Path.home() / "tinyvla_data")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--only", type=str, default=None, help="substring filter on repo_id")
    args = parser.parse_args()

    args.root.mkdir(parents=True, exist_ok=True)

    from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata

    for repo_id, n_eps in SUBSETS:
        if args.only and args.only not in repo_id:
            continue
        print(f"\n=== {repo_id} (episodes: {n_eps or 'all'}) ===")
        meta = LeRobotDatasetMetadata(repo_id)
        total_eps = meta.total_episodes
        print(f"total episodes: {total_eps}, fps: {meta.fps}")

        episodes = list(range(min(n_eps, total_eps))) if n_eps else None

        # probe with 1% to estimate size
        probe_n = max(10, (len(episodes) if episodes else total_eps) // 100)
        probe_dir = args.root / "_probe" / repo_id.replace("/", "__")
        if not args.dry_run and not probe_dir.exists():
            LeRobotDataset(repo_id, root=probe_dir, episodes=list(range(probe_n)))
            probe_gb = used_gb(probe_dir)
            est_gb = probe_gb / probe_n * (len(episodes) if episodes else total_eps)
            print(f"probe: {probe_n} eps = {probe_gb:.2f}GB -> estimated full subset {est_gb:.0f}GB")
            if est_gb > free_gb(args.root) - 100 or used_gb(args.root) + est_gb > DISK_BUDGET_GB:
                print(f"SKIP {repo_id}: would exceed budget ({DISK_BUDGET_GB}GB) or disk")
                continue

        if args.dry_run:
            continue

        target = args.root / repo_id.replace("/", "__")
        LeRobotDataset(repo_id, root=target, episodes=episodes)
        print(f"downloaded {repo_id}: {used_gb(target):.1f}GB, total used {used_gb(args.root):.1f}GB")


if __name__ == "__main__":
    main()