"""
Download TemporalBench (full, ~59 GB) and a subset of LLaVA-Video-178K (annotations
+ academic_v0_1 video splits, ~250 GB). Resumable: re-running picks up where it
left off (HF snapshot_download deduplicates).

Run via: nohup ... python download_datasets.py > log 2>&1 &
"""
from __future__ import annotations

import argparse
import os
import sys
import time
from pathlib import Path

from huggingface_hub import snapshot_download

ROOT = Path("/mnt/local-fast/opd_zt/data/raw")

# What to pull from LLaVA-Video-178K. Annotations are ~2 GB, academic splits are
# the cleaner half. We skip youtube splits (~950 GB) since 4k samples don't need
# that much.
LV178K_ALLOW = [
    "*.json",
    "*.md",
    "0_30_s_academic_v0_1/*",
    "30_60_s_academic_v0_1/*",
    "1_2_m_academic_v0_1/*",
    "2_3_m_academic_v0_1/*",
    "0_30_s_perceptiontest/*",
    "30_60_s_perceptiontest/*",
    "0_30_s_nextqa/*",
    "30_60_s_nextqa/*",
    "1_2_m_nextqa/*",
    "2_3_m_nextqa/*",
    "0_30_s_activitynetqa/*",
    "30_60_s_activitynetqa/*",
    "1_2_m_activitynetqa/*",
    "2_3_m_activitynetqa/*",
]


def log(msg: str) -> None:
    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {msg}", flush=True)


def pull(repo_id: str, local_subdir: str, allow_patterns=None) -> None:
    target = ROOT / local_subdir
    target.mkdir(parents=True, exist_ok=True)
    log(f"START  {repo_id} -> {target}  allow={allow_patterns!r}")
    t0 = time.time()
    snapshot_download(
        repo_id=repo_id,
        repo_type="dataset",
        local_dir=str(target),
        allow_patterns=allow_patterns,
        max_workers=8,
        # Resumable by default; HF skips already-downloaded files.
    )
    dt = time.time() - t0
    log(f"DONE   {repo_id}  in {dt/60:.1f} min")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--only",
        choices=["tb", "lv178k_anno", "lv178k_videos", "all"],
        default="all",
    )
    args = parser.parse_args()

    log(f"start.  ROOT={ROOT}  HF_HOME={os.environ.get('HF_HOME')}")

    if args.only in ("all", "tb"):
        pull("microsoft/TemporalBench", "tb")

    if args.only in ("all", "lv178k_anno"):
        # Annotations first so we can start filtering even before videos finish.
        pull(
            "lmms-lab/LLaVA-Video-178K",
            "lv178k",
            allow_patterns=["*.json", "*.md", "*.py"],
        )

    if args.only in ("all", "lv178k_videos"):
        pull("lmms-lab/LLaVA-Video-178K", "lv178k", allow_patterns=LV178K_ALLOW)

    log("ALL DONE")


if __name__ == "__main__":
    try:
        main()
    except Exception:
        import traceback
        traceback.print_exc()
        sys.exit(1)