"""
Downloads all required models into ./models/
Run this AFTER setup.bat

Models:
  1. Qwen2.5-VL-7B Q4_K_M  (~4.68 GB) — vision + language brain
  2. Qwen2.5-VL mmproj      (~300 MB)  — lets AI "see" the image
  Whisper + Coqui TTS auto-download on first app.py run
"""

from pathlib import Path
from huggingface_hub import hf_hub_download, list_repo_files
import os

MODELS_DIR = Path("./models")
MODELS_DIR.mkdir(exist_ok=True)

def download_file(repo_id, filename, desc):
    print(f"\n⬇️  {desc}")
    print(f"   repo : {repo_id}")
    print(f"   file : {filename}")
    print(f"   to   : {MODELS_DIR}/")
    print("   (progress bar will appear below...)\n")
    path = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        local_dir=str(MODELS_DIR),
    )
    size_gb = os.path.getsize(path) / (1024**3)
    print(f"\n   ✅ Saved → {path}  ({size_gb:.2f} GB)")
    return path

print("\n=== SunoDoc Model Downloader ===")
print("Repo   : unsloth/Qwen2.5-VL-7B-Instruct-GGUF  (confirmed correct)")
print("Total  : ~5 GB one-time download\n")

# ── 1. Find the mmproj filename dynamically ───────────────────────────────────
print("🔍 Scanning repo for mmproj file...")
REPO = "unsloth/Qwen2.5-VL-7B-Instruct-GGUF"
all_files = list(list_repo_files(REPO))
mmproj_files = [f for f in all_files if "mmproj" in f.lower()]
q4_files     = [f for f in all_files if "q4_k_m" in f.lower() and f.endswith(".gguf")]

print(f"   Found Q4_K_M files : {q4_files}")
print(f"   Found mmproj files : {mmproj_files}")

if not q4_files:
    print("\n❌ Q4_K_M file not found in repo. Available files:")
    for f in all_files:
        print(f"   {f}")
    exit(1)

MAIN_FILE   = q4_files[0]
MMPROJ_FILE = mmproj_files[0] if mmproj_files else None

# ── 2. Download main model ────────────────────────────────────────────────────
download_file(REPO, MAIN_FILE, "Qwen2.5-VL-7B main model (~4.68 GB) — the AI brain")

# ── 3. Download mmproj (vision connector) ────────────────────────────────────
if MMPROJ_FILE:
    download_file(REPO, MMPROJ_FILE, "mmproj vision connector (~300 MB)")
else:
    print("\n⚠️  No mmproj file found in this repo.")
    print("   Trying ggml-org/Qwen2.5-VL-7B-Instruct-GGUF as fallback...")
    FALLBACK_REPO = "ggml-org/Qwen2.5-VL-7B-Instruct-GGUF"
    fb_files = list(list_repo_files(FALLBACK_REPO))
    fb_mmproj = [f for f in fb_files if "mmproj" in f.lower()]
    if fb_mmproj:
        download_file(FALLBACK_REPO, fb_mmproj[0], "mmproj vision connector (fallback repo)")
    else:
        print("   ❌ mmproj not found in fallback either. Check HuggingFace manually.")

# ── 4. Summary ────────────────────────────────────────────────────────────────
print("\n=== Download complete! ===\n")
print("Files in ./models/:")
for f in sorted(MODELS_DIR.iterdir()):
    if f.is_file():
        size_gb = f.stat().st_size / (1024**3)
        print(f"  {f.name:60s}  {size_gb:.2f} GB")

print("\n📥 Whisper (STT) will auto-download on first app.py run (~500 MB)")
print("📥 Coqui Hindi TTS will auto-download on first TTS call (~100 MB)")
print("\n✅ Next step: python app.py")