"""
models.py — Descarga verificada de los modelos (DeepFilterNet3, DNSMOS y Silero VAD).

Cada fichero tiene una fuente principal (repositorio original en GitHub) y una
réplica en Hugging Face fijada a un commit concreto. Se comprueba el SHA-256
siempre que se conoce, así que da igual de dónde venga: o es el fichero bueno
o no se usa.
"""

from __future__ import annotations

import hashlib
import os
import shutil
import tempfile
import time
import urllib.request
import zipfile
from pathlib import Path

CACHE_DIR = Path(os.environ.get("AUDIO_DENOISE_CACHE", Path.home() / ".cache" / "audio-denoise"))

_HF_DFN = "https://huggingface.co/niobures/DeepFilterNet/resolve/563a7658343bc716f2bca304f16e54f403914ce1"
_HF_DNSMOS = "https://huggingface.co/anchor-flux/dnsmos-onnx/resolve/4323f201af8a33969025d4e39b409205f0881237"
_GH_DNSMOS = "https://github.com/microsoft/DNS-Challenge/raw/master/DNSMOS/DNSMOS"

DFN3_ZIP = {
    "name": "DeepFilterNet3.zip",
    "sha256": "49c52edc8947ae1f9bf50d81530beaf3a2c3245aeaf34b6f31ff535cd22284d2",
    "urls": [
        "https://github.com/Rikorose/DeepFilterNet/raw/main/models/DeepFilterNet3.zip",
        f"{_HF_DFN}/models/onnx/DeepFilterNet/DeepFilterNet3.zip",
    ],
}

DNSMOS_FILES = {
    "sig_bak_ovr.onnx": {
        "sha256": "269fbebdb513aa23cddfbb593542ecc540284a91849ac50516870e1ac78f6edd",
        "urls": [f"{_GH_DNSMOS}/sig_bak_ovr.onnx", f"{_HF_DNSMOS}/sig_bak_ovr.onnx"],
    },
}


_HF_SILERO = "https://huggingface.co/onnx-community/silero-vad/resolve/e71cae966052b992a7eca6b17738916ce0eca4ec"
SILERO_VAD = {
    "name": "silero_vad.onnx",
    "sha256": "a4a068cd6cf1ea8355b84327595838ca748ec29a25bc91fc82e6c299ccdc5808",
    "urls": [f"{_HF_SILERO}/onnx/model.onnx"],
}


def _sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()


def _download(urls: list[str], dest: Path, sha256: str, retries: int = 3) -> Path:
    """Descarga `dest` probando cada URL; valida el hash. Lanza RuntimeError si todo falla."""
    if dest.exists() and _sha256(dest) == sha256:
        return dest
    dest.parent.mkdir(parents=True, exist_ok=True)
    errors = []
    for url in urls:
        for attempt in range(retries):
            try:
                with tempfile.NamedTemporaryFile(dir=dest.parent, delete=False) as tmp:
                    req = urllib.request.Request(url, headers={"User-Agent": "audio-denoise"})
                    with urllib.request.urlopen(req, timeout=60) as r:
                        shutil.copyfileobj(r, tmp)
                tmp_path = Path(tmp.name)
                digest = _sha256(tmp_path)
                if digest != sha256:
                    tmp_path.unlink(missing_ok=True)
                    errors.append(f"{url}: hash distinto ({digest[:12]}…)")
                    break  # reintentar la misma URL no arregla un hash distinto
                tmp_path.replace(dest)
                print(f"[MODELOS] {dest.name} ← {url}")
                return dest
            except Exception as e:  # red, 403, timeout…
                errors.append(f"{url}: {e}")
                time.sleep(2 ** attempt)
    raise RuntimeError(f"No se pudo descargar {dest.name}:\n" + "\n".join(errors))


def dfn3_dir() -> Path:
    """Devuelve la carpeta del modelo DeepFilterNet3 (config.ini + checkpoints/)."""
    model_dir = CACHE_DIR / "DeepFilterNet3"
    if (model_dir / "config.ini").is_file() and (model_dir / "checkpoints").is_dir():
        return model_dir
    zip_path = _download(DFN3_ZIP["urls"], CACHE_DIR / DFN3_ZIP["name"], DFN3_ZIP["sha256"])
    with zipfile.ZipFile(zip_path) as z:
        z.extractall(CACHE_DIR)
    return model_dir


def dnsmos_paths() -> dict[str, Path]:
    """Devuelve {nombre: ruta} de los modelos ONNX de DNSMOS."""
    return {
        name: _download(spec["urls"], CACHE_DIR / "dnsmos" / name, spec["sha256"])
        for name, spec in DNSMOS_FILES.items()
    }


def silero_vad_path() -> Path:
    """Detector de voz Silero VAD v5 (ONNX, licencia MIT)."""
    return _download(SILERO_VAD["urls"], CACHE_DIR / "silero" / SILERO_VAD["name"], SILERO_VAD["sha256"])