#!/usr/bin/env python
"""Push a packed wds dataset dir (shard-*.tar + manifest.json + tasks.json) to
the Hub as a private dataset repo, and verify the listing afterwards.

Usage:
    python hf_push_wds.py --src /data/wds/tartandrive --repo AlexWortega/wds-tartandrive
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", type=Path, required=True)
    ap.add_argument("--repo", required=True)
    ap.add_argument("--public", action="store_true")
    args = ap.parse_args()

    from huggingface_hub import HfApi

    # single packed dataset (manifest at top) or a bundle of per-task subdirs
    if (args.src / "manifest.json").exists():
        dirs = {"": json.loads((args.src / "manifest.json").read_text())}
    else:
        dirs = {p.name + "/": json.loads((p / "manifest.json").read_text())
                for p in sorted(args.src.iterdir())
                if p.is_dir() and (p / "manifest.json").exists()}
        if not dirs:
            raise SystemExit(f"{args.src}: no manifest.json anywhere")

    api = HfApi()
    api.create_repo(args.repo, repo_type="dataset", private=not args.public, exist_ok=True)
    api.upload_large_folder(repo_id=args.repo, repo_type="dataset",
                            folder_path=str(args.src))

    remote = set(api.list_repo_files(args.repo, repo_type="dataset"))
    missing, total = [], 0
    for pfx, mf in dirs.items():
        missing += [pfx + s for s in mf["shards"] if pfx + s not in remote]
        missing += [pfx + f for f in ("manifest.json", "tasks.json") if pfx + f not in remote]
        total += mf["samples"]
    if missing:
        raise SystemExit(f"upload incomplete, missing: {missing[:10]}")
    print(f"{args.repo}: {len(dirs)} dataset dir(s), {total:,} samples, all shards verified")


if __name__ == "__main__":
    main()