#!/usr/bin/env python
"""Fetch wds dataset repos (hf_push_wds.py output) to a local root for training.

Reads the same specs yaml train_fast.py uses (source.type: wds): every entry
with an hf_repo is downloaded to <root>/<repo basename> unless already complete
(all shards from manifest.json present).

Usage:
    python fetch_wds.py --specs configs/wds_mix.yaml --root /data/wds
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import yaml


def _complete(d: Path) -> bool:
    mf_local = d / "manifest.json"
    if not mf_local.exists():
        return False
    mf = json.loads(mf_local.read_text())
    return all((d / s).exists() for s in mf["shards"])


def _dataset_dirs(root: Path):
    """A repo is either one packed dataset (manifest at top) or a bundle of
    per-task subdirs each with its own manifest (wds-unitree)."""
    if (root / "manifest.json").exists():
        return [root]
    return [p for p in sorted(root.iterdir()) if p.is_dir() and (p / "manifest.json").exists()]


def fetch(repo: str, dest: Path):
    from huggingface_hub import snapshot_download

    dest.mkdir(parents=True, exist_ok=True)
    dirs = _dataset_dirs(dest)
    if dirs and all(_complete(d) for d in dirs):
        print(f"{repo}: already complete ({len(dirs)} dataset dir(s))")
        return
    snapshot_download(repo, repo_type="dataset", local_dir=str(dest))
    dirs = _dataset_dirs(dest)
    if not dirs:
        raise SystemExit(f"{repo}: no manifest.json found after download")
    for d in dirs:
        if not _complete(d):
            raise SystemExit(f"{repo}: incomplete after download: {d}")
        mf = json.loads((d / "manifest.json").read_text())
        print(f"{repo}::{d.name}: {len(mf['shards'])} shards, {mf['samples']:,} samples")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--specs", type=Path, required=True)
    ap.add_argument("--root", type=Path, required=True)
    args = ap.parse_args()
    specs = yaml.safe_load(args.specs.read_text())["datasets"]
    done = set()
    for s in specs:
        repo = s.get("hf_repo")
        if repo and "dir" not in s and repo not in done:
            fetch(repo, args.root / repo.split("/")[-1])
            done.add(repo)
        elif not repo:
            print(f"local: {s.get('dir')}")


if __name__ == "__main__":
    main()