"""
Download models from HuggingFace Hub.
Usage:
    python fetch_models.py              # Download all models
    python fetch_models.py --lora       # Re-pull latest (Nemotron) LoRA adapter
    python fetch_models.py --qwen-lora  # Pull old Qwen LoRA from specific commit
    python fetch_models.py --nemotron   # Download Nemotron base model
    python fetch_models.py --base       # Download base Qwen model only
"""
import sys
import os
from huggingface_hub import snapshot_download

# HF_TOKEN = "your_token_here"
HF_TOKEN = os.environ.get("HF_TOKEN")
LORA_REPO = "Rolaficus/Qwen_7B_Cuda"
BASE_MODEL = "unsloth/qwen2.5-coder-7b-bnb-4bit"
NEMOTRON_MODEL = "unsloth/NVIDIA-Nemotron-Nano-9B-v2"

# The old Qwen LoRA commit (before it was replaced with Nemotron LoRA)
QWEN_LORA_REVISION = "7b4d77303676"


def download_lora():
    """Download latest LoRA (Nemotron adapter) from the repo."""
    print(f"Downloading latest LoRA (Nemotron): {LORA_REPO}...")
    snapshot_download(
        repo_id=LORA_REPO,
        local_dir="Qwen_7B_Cuda",
        token=HF_TOKEN,
    )
    print("Latest (Nemotron) LoRA downloaded successfully!\n")


def download_qwen_lora():
    """Download old Qwen LoRA from a specific commit."""
    print(f"Downloading Qwen LoRA from commit {QWEN_LORA_REVISION}...")
    snapshot_download(
        repo_id=LORA_REPO,
        revision=QWEN_LORA_REVISION,
        local_dir="Qwen_7B_Cuda_Qwen",
        token=HF_TOKEN,
    )
    print("Qwen LoRA downloaded successfully!\n")


def download_base():
    print(f"Downloading Base Model: {BASE_MODEL}...")
    snapshot_download(
        repo_id=BASE_MODEL,
        local_dir="base_model",
        token=HF_TOKEN,
    )
    print("Base model weights downloaded successfully!\n")


def download_nemotron():
    print(f"Downloading Nemotron: {NEMOTRON_MODEL}...")
    print("   This is a large download (~18GB). Please be patient.")
    snapshot_download(
        repo_id=NEMOTRON_MODEL,
        local_dir="nemotron_9b",
        token=HF_TOKEN,
    )
    print("Nemotron model downloaded successfully!\n")


def download_all():
    download_lora()
    download_qwen_lora()
    download_base()
    download_nemotron()


if __name__ == "__main__":
    args = sys.argv[1:]

    if not args:
        download_all()
    else:
        if "--lora" in args:
            download_lora()
        if "--qwen-lora" in args:
            download_qwen_lora()
        if "--base" in args:
            download_base()
        if "--nemotron" in args:
            download_nemotron()

    print("All requested downloads complete!")