"""Download the required weights.

  * Qwen/Qwen3-0.6B            - System-1 backbone (NanoJev-style student)
  * Qwen/Qwen3.5-9B            - System-2 multimodal controller / teacher
  * C-Tianyu/NanoJev           - optional reference checkpoint (unified-games-v1)

Tokens: gated repos need HF_TOKEN; the NanoJev and Qwen repos are public.
"""

from __future__ import annotations

import argparse


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--skip-nanojev", action="store_true")
    args = ap.parse_args()

    from huggingface_hub import snapshot_download

    print(">> Qwen3-0.6B (System-1 backbone)")
    snapshot_download("Qwen/Qwen3-0.6B", local_dir="checkpoints/Qwen3-0.6B")

    if not args.skip_nanojev:
        print(">> NanoJev reference checkpoint (optional, for head surgery)")
        snapshot_download(
            "C-Tianyu/NanoJev",
            revision="unified-games-v1",
            local_dir="checkpoints/NanoJev-ref",
            allow_patterns=["best.safetensors", "config.json", "tokenizer/*"],
        )

    print(">> Qwen3.5-9B (System-2 / teacher) - large download")
    snapshot_download("Qwen/Qwen3.5-9B", local_dir="checkpoints/Qwen3.5-9B")
    print("done -> checkpoints/")


if __name__ == "__main__":
    main()