#!/usr/bin/env python3
"""Census of the public LeRobot dataset index. Anonymous and read-only. Reads metadata only.

  python3 fetch.py [outdir]

1. Lists every dataset tagged LeRobot through the public listing endpoint with filter=LeRobot,
   following the Link header. (other=LeRobot is ignored by that endpoint and returns untagged results.)
2. Takes a seeded random sample of up to 1,000 ids.
3. Reads meta/info.json for each, at no more than 2 requests a second. Gated or private repos are skipped
   and counted. A 429 stops the run and whatever was read is kept. No episode data is downloaded.
"""
import json
import pathlib
import random
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

LIST = "https://huggingface.co/api/datasets?filter=LeRobot&limit=1000"
INFO = "https://huggingface.co/datasets/{id}/resolve/main/meta/info.json"
SEED = 20261108
SAMPLE = 1000
MIN_GAP = 0.5  # seconds, so at most 2 requests a second


class RateLimited(Exception):
    pass


last = 0.0


def get(url):
    global last
    wait = MIN_GAP - (time.time() - last)
    if wait > 0:
        time.sleep(wait)
    last = time.time()
    req = urllib.request.Request(url)
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            return r.status, r.headers, r.read()
    except urllib.error.HTTPError as ex:
        if ex.code == 429:
            raise RateLimited(url)
        return ex.code, ex.headers, b""


def enumerate_index():
    items, url, complete, pages = [], LIST, True, 0
    while url:
        try:
            status, headers, body = get(url)
        except RateLimited:
            complete = False
            print("rate limited while listing, stopping list", file=sys.stderr)
            break
        if status != 200:
            complete = False
            print("listing stopped with status", status, file=sys.stderr)
            break
        pages += 1
        batch = json.loads(body)
        items += batch
        m = re.search(r'<([^>]+)>;\s*rel="next"', headers.get("Link", "") or "")
        url = m.group(1) if m else None
        print("page", pages, "items", len(items), file=sys.stderr)
    return items, complete, pages


def main():
    out = pathlib.Path(sys.argv[1]) if len(sys.argv) > 1 else pathlib.Path(__file__).parent
    out.mkdir(parents=True, exist_ok=True)
    items, complete, pages = enumerate_index()
    ids = sorted(x["id"] for x in items)
    by_id = {x["id"]: x for x in items}
    rng = random.Random(SEED)
    sample = rng.sample(ids, min(SAMPLE, len(ids)))
    meta = {"date": time.strftime("%Y-%m-%d"), "index_total": len(ids), "index_complete": complete, "index_pages": pages,
            "seed": SEED, "sample_requested": len(sample), "skipped_gated_or_private": 0, "no_info_json": 0, "errors": 0, "read": 0,
            "rate_limited": False}
    (out / "index_ids.json").write_text(json.dumps(ids))
    raw = (out / "raw.jsonl").open("w", encoding="utf-8")
    for i, dsid in enumerate(sample):
        x = by_id[dsid]
        if x.get("gated") or x.get("private"):
            meta["skipped_gated_or_private"] += 1
            continue
        tags = x.get("tags", [])
        licence = sorted(t.split(":", 1)[1] for t in tags if t.startswith("license:"))
        try:
            status, _, body = get(INFO.format(id=urllib.parse.quote(dsid, safe="/")))
        except RateLimited:
            meta["rate_limited"] = True
            meta["stopped_at"] = i
            print("rate limited, stopping sample read", file=sys.stderr)
            break
        if status in (401, 403):
            meta["skipped_gated_or_private"] += 1
            continue
        if status == 404:
            meta["no_info_json"] += 1
            continue
        if status != 200:
            meta["errors"] += 1
            continue
        try:
            info = json.loads(body)
        except ValueError:
            meta["errors"] += 1
            continue
        features = info.get("features") or {}
        streams = [k for k, v in features.items() if isinstance(v, dict) and v.get("dtype") in ("video", "image")]
        raw.write(json.dumps({"id": dsid, "licence": licence, "robot_type": info.get("robot_type"), "episodes": info.get("total_episodes"),
                              "fps": info.get("fps"), "cameras": len(streams), "version": info.get("codebase_version")}) + "\n")
        meta["read"] += 1
        if i % 50 == 0:
            print("sample", i, "read", meta["read"], file=sys.stderr)
    raw.close()
    (out / "meta.json").write_text(json.dumps(meta, indent=1))
    print(json.dumps(meta))


if __name__ == "__main__":
    main()