#!/usr/bin/env python
"""Prepare the SO-100/101 slice of HuggingFaceVLA/community_dataset_v3.

Steps:
  catalog  — fetch all sub-dataset info.json, write catalog.json
  select   — filter SO100/101 single-arm, rank by episodes, write selection.json
  convert  — download each selected subdir + convert v2.1 -> v3.0 locally

Usage:
    python scripts/prepare_community_v3.py catalog
    python scripts/prepare_community_v3.py select --target-episodes 12000
    python scripts/prepare_community_v3.py convert [--limit N]
"""

from __future__ import annotations

import argparse
import json
import shutil
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

AGG_REPO = "HuggingFaceVLA/community_dataset_v3"
DISK_BUDGET_GB = 400
DATA_ROOT = Path.home() / "tinyvla_data"
CATALOG = DATA_ROOT / "community_v3_catalog.json"
SELECTION = DATA_ROOT / "community_v3_selection.json"
CONVERTED_DIR = DATA_ROOT / "so101_v3"
SO_ROBOTS = {"so100", "so101", "so100_follower", "so101_follower"}


def cmd_catalog(args):
    from huggingface_hub import HfApi, hf_hub_download

    api = HfApi()
    files = api.list_repo_files(AGG_REPO, repo_type="dataset")
    subs = sorted(f.rsplit("/meta/info.json", 1)[0] for f in files if f.endswith("meta/info.json"))
    print(f"{len(subs)} sub-datasets")

    def fetch(sub):
        try:
            p = hf_hub_download(AGG_REPO, f"{sub}/meta/info.json", repo_type="dataset")
            info = json.load(open(p))
            return sub, {
                "robot_type": info.get("robot_type"),
                "episodes": info.get("total_episodes"),
                "frames": info.get("total_frames"),
                "fps": info.get("fps"),
                "version": info.get("codebase_version"),
                "action_shape": info.get("features", {}).get("action", {}).get("shape"),
                "cameras": [k for k in info.get("features", {}) if k.startswith("observation.images")],
            }
        except Exception as e:
            return sub, {"error": str(e)[:100]}

    catalog = {}
    with ThreadPoolExecutor(16) as ex:
        futs = [ex.submit(fetch, s) for s in subs]
        for i, f in enumerate(as_completed(futs)):
            sub, meta = f.result()
            catalog[sub] = meta
            if (i + 1) % 100 == 0:
                print(f"{i+1}/{len(subs)}")

    DATA_ROOT.mkdir(parents=True, exist_ok=True)
    CATALOG.write_text(json.dumps(catalog, indent=1))
    print(f"wrote {CATALOG}")


def cmd_select(args):
    catalog = json.loads(CATALOG.read_text())
    rows = [
        (sub, m)
        for sub, m in catalog.items()
        if m.get("robot_type") in SO_ROBOTS
        and m.get("action_shape") == [6]
        and m.get("fps") == 30
        and m.get("episodes")
        and m.get("cameras")
    ]
    # prefer larger datasets: fewer conversions per episode
    rows.sort(key=lambda r: -r[1]["episodes"])
    picked, total = [], 0
    for sub, m in rows:
        if total >= args.target_episodes:
            break
        picked.append({"sub": sub, **m})
        total += m["episodes"]
    print(f"{len(rows)} eligible; picked {len(picked)} datasets, {total} episodes")
    SELECTION.write_text(json.dumps(picked, indent=1))
    print(f"wrote {SELECTION}")


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


def used_gb_cached(path: Path) -> float:
    import subprocess

    out = subprocess.run(["du", "-s", "--block-size=1G", str(path)], capture_output=True, text=True)
    return float(out.stdout.split()[0]) if out.returncode == 0 else 0.0


def cmd_convert(args):
    from huggingface_hub import snapshot_download
    from lerobot.scripts.convert_dataset_v21_to_v30 import convert_dataset

    picked = json.loads(SELECTION.read_text())
    if args.limit:
        picked = picked[: args.limit]
    CONVERTED_DIR.mkdir(parents=True, exist_ok=True)
    raw_dir = DATA_ROOT / "_raw_v21"
    done, failed = 0, []
    for item in picked:
        sub = item["sub"]
        name = sub.replace("/", "__")
        out = CONVERTED_DIR / name
        if (out / "meta" / "info.json").exists():
            done += 1
            continue
        local = raw_dir / name
        try:
            snapshot_download(
                AGG_REPO,
                repo_type="dataset",
                allow_patterns=[f"{sub}/*"],
                local_dir=raw_dir / "_dl",
            )
            src = raw_dir / "_dl" / sub
            if local.exists():
                shutil.rmtree(local)
            shutil.move(str(src), str(local))
            convert_dataset(repo_id=name, root=local, push_to_hub=False, force_conversion=True)
            # converter writes v3.0 in place at root
            shutil.move(str(local), str(out))
            # drop the v2.1 originals the converter stashes as <root>_old
            old = local.parent / (local.name + "_old")
            if old.exists():
                shutil.rmtree(old)
            done += 1
            print(f"[{done}/{len(picked)}] {sub}: converted -> {out} (free {free_gb(DATA_ROOT):.0f}GB)")
            if free_gb(DATA_ROOT) < 150 or used_gb_cached(DATA_ROOT) > DISK_BUDGET_GB:
                print("STOP: disk budget reached")
                break
        except Exception as e:
            failed.append(sub)
            print(f"FAIL {sub}: {type(e).__name__}: {str(e)[:200]}")
    print(f"done={done} failed={len(failed)}")
    if failed:
        (DATA_ROOT / "convert_failures.json").write_text(json.dumps(failed))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="cmd", required=True)
    sub.add_parser("catalog")
    p_sel = sub.add_parser("select")
    p_sel.add_argument("--target-episodes", type=int, default=12_000)
    p_conv = sub.add_parser("convert")
    p_conv.add_argument("--limit", type=int, default=None)
    args = parser.parse_args()
    {"catalog": cmd_catalog, "select": cmd_select, "convert": cmd_convert}[args.cmd](args)