"""Fetch + extract the WISA-80K wan480p shards.

The clips live in ``csusupergear/WISA-80K-wan480p-16fps-81f`` as 276 webdataset
tars; shard ``N`` holds ids ``N*100 .. N*100+99`` as ``<id>.{mp4,txt,json}``.
Download and extraction are interleaved over a thread pool so a shard is unpacked
as soon as it lands.  Both steps are idempotent: an already-downloaded shard is
served from the HF cache and an already-extracted shard is skipped.
"""

from __future__ import annotations

import argparse
import os
import tarfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

REPO = "csusupergear/WISA-80K-wan480p-16fps-81f"


def shard_done(vdir: Path, shard: int) -> bool:
    """A shard counts as extracted when its last clip's three files exist."""
    last = f"{shard * 100 + 99:08d}"
    return all((vdir / f"{last}.{ext}").exists() for ext in ("mp4", "txt", "json"))


def fetch_one(shard: int, raw_dir: Path, vdir: Path) -> tuple[int, str]:
    from huggingface_hub import hf_hub_download

    if shard_done(vdir, shard):
        return shard, "skip"
    path = hf_hub_download(
        REPO, f"{shard:06d}.tar", repo_type="dataset",
        local_dir=str(raw_dir),
    )
    with tarfile.open(path) as tf:
        tf.extractall(vdir)
    return shard, "ok"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--num_shards", type=int, default=276)
    ap.add_argument("--raw_dir", default="data/raw/wisa_wan480p")
    ap.add_argument("--video_dir", default="data/videos")
    ap.add_argument("--workers", type=int, default=8)
    args = ap.parse_args()

    raw_dir, vdir = Path(args.raw_dir), Path(args.video_dir)
    vdir.mkdir(parents=True, exist_ok=True)

    done = n_err = 0
    with ThreadPoolExecutor(args.workers) as ex:
        futs = {ex.submit(fetch_one, s, raw_dir, vdir): s
                for s in range(args.num_shards)}
        for f in as_completed(futs):
            s = futs[f]
            try:
                _, status = f.result()
            except Exception as e:                       # network / corrupt tar
                n_err += 1
                print(f"[err] shard {s}: {e}", flush=True)
                continue
            done += 1
            if done % 10 == 0:
                print(f"{done}/{args.num_shards} shards ({status} {s:06d})",
                      flush=True)

    n = len(list(vdir.glob("*.mp4")))
    print(f"DONE shards={done} errors={n_err} clips={n}", flush=True)


if __name__ == "__main__":
    main()