"""
download_models.py
──────────────────
Pre-downloads the two inference models used by Aara so the first
request is fast. Called once at Docker build time (or run manually).

Models:
  1. faster-whisper-medium (~769 MB, better CPU accuracy for hotel speech)
  2. Qwen2.5-0.5B-Instruct Q4_K_M GGUF (~330 MB, ~15-25 tok/s on CPU)
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

MODELS_DIR = Path(os.environ.get("AARA_MODELS_DIR", "models"))
HF_HOME = Path(os.environ.get("HF_HOME", "cache/hf"))


def _whisper_dirname(repo_id: str) -> str:
    override = os.environ.get("AARA_WHISPER_MODEL_DIRNAME", "").strip()
    if override:
        return override
    name = repo_id.rsplit("/", 1)[-1].strip().lower()
    if name.startswith("faster-whisper-"):
        return "whisper-" + name[len("faster-whisper-"):]
    cleaned = name.replace("faster-", "")
    if cleaned.startswith("whisper-"):
        return cleaned
    if cleaned.startswith("distil-") or cleaned.endswith("-turbo") or "large-v3" in cleaned or cleaned in {"small", "base", "medium"}:
        return "whisper-" + cleaned
    return cleaned


WHISPER_REPO = os.environ.get("AARA_WHISPER_MODEL", "Systran/faster-whisper-medium")
WHISPER_LOCAL = MODELS_DIR / _whisper_dirname(WHISPER_REPO)

QWEN_REPO = "Qwen/Qwen2.5-0.5B-Instruct-GGUF"
QWEN_FILENAME = "qwen2.5-0.5b-instruct-q4_k_m.gguf"
QWEN_LOCAL = MODELS_DIR / "qwen-gguf"


def _ensure_dirs() -> None:
    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    QWEN_LOCAL.mkdir(parents=True, exist_ok=True)
    HF_HOME.mkdir(parents=True, exist_ok=True)


def download_whisper() -> None:
    target = WHISPER_LOCAL / "config.json"
    if target.exists():
        print(f"[skip] Whisper already at {WHISPER_LOCAL}")
        return

    print(f"[download] {WHISPER_REPO} → {WHISPER_LOCAL} ...")
    from huggingface_hub import snapshot_download

    snapshot_download(
        repo_id=WHISPER_REPO,
        local_dir=str(WHISPER_LOCAL),
        local_dir_use_symlinks=False,
    )
    print(f"[ok] Whisper model downloaded: {WHISPER_REPO}")


def download_qwen_gguf() -> None:
    target = QWEN_LOCAL / QWEN_FILENAME
    if target.exists():
        print(f"[skip] Qwen GGUF already at {target}")
        return

    print(f"[download] {QWEN_REPO}/{QWEN_FILENAME} → {QWEN_LOCAL} ...")
    from huggingface_hub import hf_hub_download

    hf_hub_download(
        repo_id=QWEN_REPO,
        filename=QWEN_FILENAME,
        local_dir=str(QWEN_LOCAL),
        local_dir_use_symlinks=False,
    )
    print("[ok] Qwen Q4_K_M GGUF downloaded.")


def main() -> None:
    _ensure_dirs()
    try:
        download_whisper()
    except Exception as exc:
        print(f"[warn] Whisper download failed (will retry at runtime): {exc}", file=sys.stderr)

    try:
        download_qwen_gguf()
    except Exception as exc:
        print(f"[warn] Qwen GGUF download failed (will retry at runtime): {exc}", file=sys.stderr)

    print("[done] Model pre-download complete.")


if __name__ == "__main__":
    main()