"""Download GGUF model files for the VLM shootout.

Each VLM needs two files:
  1. The main model weights (quantized GGUF)
  2. The multimodal projector (mmproj) that maps image embeddings
     into the language model's embedding space.

We download Q4_K_M quantizations to balance quality and VRAM usage
on an 8GB GPU.
"""

from huggingface_hub import hf_hub_download
from pathlib import Path

MODELS_DIR = Path(__file__).parent.parent / "models"

MODELS = {
    "qwen2.5-vl-3b": {
        "repo": "mradermacher/Qwen2.5-VL-3B-Instruct-GGUF",
        "model_file": "Qwen2.5-VL-3B-Instruct.Q4_K_M.gguf",
        "mmproj_file": "Qwen2.5-VL-3B-Instruct.mmproj-fp16.gguf",
    },
    "smolvlm-2b": {
        "repo": "ggml-org/SmolVLM-Instruct-GGUF",
        "model_file": "SmolVLM-Instruct-Q4_K_M.gguf",
        "mmproj_file": "mmproj-SmolVLM-Instruct-f16.gguf",
    },
    "gemma-3-4b": {
        "repo": "ggml-org/gemma-3-4b-it-GGUF",
        "model_file": "gemma-3-4b-it-Q4_K_M.gguf",
        "mmproj_file": "mmproj-model-f16.gguf",
    },
}


def download_model(name: str, info: dict) -> dict[str, Path]:
    """Download model + mmproj files, return local paths."""
    print(f"\n{'='*60}")
    print(f"Downloading: {name}")
    print(f"  Repo: {info['repo']}")
    print(f"{'='*60}")

    model_path = Path(hf_hub_download(
        repo_id=info["repo"],
        filename=info["model_file"],
        local_dir=MODELS_DIR / name,
    ))
    print(f"  Model: {model_path} ({model_path.stat().st_size / 1e9:.2f} GB)")

    mmproj_path = Path(hf_hub_download(
        repo_id=info["repo"],
        filename=info["mmproj_file"],
        local_dir=MODELS_DIR / name,
    ))
    print(f"  Mmproj: {mmproj_path} ({mmproj_path.stat().st_size / 1e9:.2f} GB)")

    return {"model": model_path, "mmproj": mmproj_path}


def main():
    MODELS_DIR.mkdir(parents=True, exist_ok=True)

    import argparse
    parser = argparse.ArgumentParser(description="Download VLM GGUF models")
    parser.add_argument(
        "--model",
        choices=list(MODELS.keys()) + ["all"],
        default="all",
        help="Which model to download (default: all)",
    )
    args = parser.parse_args()

    targets = MODELS if args.model == "all" else {args.model: MODELS[args.model]}

    paths = {}
    for name, info in targets.items():
        paths[name] = download_model(name, info)

    print(f"\n{'='*60}")
    print("Download complete!")
    for name, p in paths.items():
        print(f"  {name}:")
        print(f"    model:  {p['model']}")
        print(f"    mmproj: {p['mmproj']}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()