#!/usr/bin/env python3
"""
Download and verify model weights from HuggingFace.

Usage
-----
python scripts/setup_models.py                      # all three models
python scripts/setup_models.py --model world_model  # just one
python scripts/setup_models.py --verify-only        # check what's cached
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

MODEL_MAP = {
    "world_model": "nvidia/Cosmos-Predict2-2B",
    "policy":      "nvidia/Cosmos-Policy-2B",
    "language":    "openvla/openvla-7b",
}


def download(model_key: str, repo_id: str):
    from huggingface_hub import snapshot_download

    print(f"\n{'─'*60}")
    print(f"  Downloading: {repo_id}")
    print(f"{'─'*60}")
    local = snapshot_download(repo_id=repo_id, ignore_patterns=["*.bin"])
    print(f"  Saved to: {local}")


def verify(repo_id: str) -> bool:
    from huggingface_hub import try_to_load_from_cache, HfFileSystem

    fs = HfFileSystem()
    try:
        files = fs.ls(repo_id, detail=False)
        cached = any(
            try_to_load_from_cache(repo_id, Path(f).name) for f in files
        )
        return cached
    except Exception:
        return False


def main():
    parser = argparse.ArgumentParser(description="Download Cosmos + OpenVLA weights.")
    parser.add_argument("--model", choices=list(MODEL_MAP.keys()), default=None,
                        help="Download a specific model; omit for all")
    parser.add_argument("--verify-only", action="store_true",
                        help="Print cache status without downloading")
    args = parser.parse_args()

    targets = {args.model: MODEL_MAP[args.model]} if args.model else MODEL_MAP

    for key, repo_id in targets.items():
        cached = verify(repo_id)
        status = "CACHED ✓" if cached else "not cached"
        print(f"  {key:15s}  {repo_id:45s}  {status}")

        if not args.verify_only and not cached:
            download(key, repo_id)

    print("\nDone.")


if __name__ == "__main__":
    main()