from __future__ import annotations

import os
import shutil

from huggingface_hub import hf_hub_download

from config import CONFIG, MODELS_DIR

_PIPER_REPO = "rhasspy/piper-voices"


def _fetch(repo: str, filename: str, dst_name: str) -> None:
    """Download repo/filename from the Hub into models/dst_name if missing."""
    dst = os.path.join(MODELS_DIR, dst_name)
    if os.path.exists(dst):
        return
    os.makedirs(MODELS_DIR, exist_ok=True)
    src = hf_hub_download(repo, filename)
    shutil.copy(src, dst)


def _piper_repo_path(flat: str) -> str:
    """es_MX-claude-high.onnx -> es/es_MX/claude/high/es_MX-claude-high.onnx."""
    name = flat[:-len(".onnx")] if flat.endswith(".onnx") else flat
    locale, speaker, quality = name.split("-", 2)
    lang = locale.split("_")[0]
    return f"{lang}/{locale}/{speaker}/{quality}/{name}.onnx"


def ensure_models() -> None:
    """Make sure the active backend's GGUF, LoRA, and Piper voice are present."""
    b = CONFIG.llm
    if b.hf_repo:
        _fetch(b.hf_repo, b.gguf_path, b.gguf_path)
    if b.lora_path and b.lora_repo:
        _fetch(b.lora_repo, b.lora_path, b.lora_path)

    # Piper voice (.onnx + its .onnx.json), flattened into models/.
    voice = CONFIG.piper_voice
    if not os.path.exists(CONFIG.piper_path()):
        repo_path = _piper_repo_path(voice)
        _fetch(_PIPER_REPO, repo_path, voice)
        _fetch(_PIPER_REPO, repo_path + ".json", voice + ".json")