#!/usr/bin/env python
"""Enumerate unitreerobotics LeRobot datasets and emit a conversion worklist.

Filters to repos whose meta/info.json has flat observation.state + action
features (the v2.1 G1_Dex1_* sets split state into 7 named fields — those need
a concat converter and are excluded here). Output: JSON lines sorted by frames,
ready to drive build_shards(.py|_v2.py) + pack_wds.py.

Usage:
    python list_unitree_specs.py [--min-frames 10000] > unitree_worklist.jsonl
"""

from __future__ import annotations

import argparse
import json
import sys

import requests


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--org", default="unitreerobotics")
    ap.add_argument("--min-frames", type=int, default=0)
    args = ap.parse_args()

    ds = requests.get(
        f"https://huggingface.co/api/datasets?author={args.org}&limit=200", timeout=30
    ).json()
    rows = []
    for d in ds:
        rid = d["id"]
        try:
            info = requests.get(
                f"https://huggingface.co/datasets/{rid}/resolve/main/meta/info.json",
                timeout=30, allow_redirects=True,
            ).json()
        except Exception:
            continue
        feats = info.get("features", {})
        if "observation.state" not in feats or "action" not in feats:
            print(f"skip {rid}: no flat state/action", file=sys.stderr)
            continue
        cams = [k for k, v in feats.items() if v.get("dtype") == "video"]
        rows.append({
            "repo_id": rid,
            "version": info["codebase_version"],
            "fps": info["fps"],
            "episodes": info["total_episodes"],
            "frames": info["total_frames"],
            "state_dim": feats["observation.state"]["shape"][0],
            "action_dim": feats["action"]["shape"][0],
            "cameras": cams,
        })
    rows.sort(key=lambda r: -r["frames"])
    kept = [r for r in rows if r["frames"] >= args.min_frames]
    for r in kept:
        print(json.dumps(r))
    tot = sum(r["frames"] for r in kept)
    print(f"{len(kept)} datasets, {tot:,} frames total", file=sys.stderr)


if __name__ == "__main__":
    main()