"""
denoise_core.py — Motor de reducción de ruido (DeepFilterNet3) con ZeroGPU y respaldo en CPU.

- El modelo trabaja a 48 kHz. Quien llama se encarga de remuestrear (ver dsp.py).
- El audio se procesa por bloques con contexto y fundido cruzado, para que cada
  llamada a la GPU sea corta y la memoria no dependa de la duración del fichero.
- Si la GPU no está disponible (cuota agotada, cola, error), se usa la CPU.
"""

from __future__ import annotations

import copy
import sys
import time
import types
import warnings

try:
    import spaces  # en ZeroGPU debe importarse antes que torch
except ImportError:  # versión de escritorio: sin ZeroGPU
    spaces = None

# Shim de compatibilidad: deepfilternet 0.5.6 importa torchaudio.backend.common (solo para df.io,
# que no usamos). Versiones recientes de torchaudio lo han eliminado: si falta, se crea uno mínimo.
try:
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        import torchaudio.backend.common  # noqa: F401
except ImportError:
    from dataclasses import dataclass

    @dataclass
    class AudioMetaData:
        sample_rate: int
        num_frames: int
        num_channels: int
        bits_per_sample: int
        encoding: str

    _common = types.ModuleType("torchaudio.backend.common")
    _common.AudioMetaData = AudioMetaData
    _backend = types.ModuleType("torchaudio.backend")
    _backend.common = _common
    sys.modules["torchaudio.backend"] = _backend
    sys.modules["torchaudio.backend.common"] = _common

import numpy as np
import torch
from torch.nn import functional as F

with warnings.catch_warnings():
    warnings.simplefilter("ignore")  # avisos de deprecación de torchaudio dentro de deepfilternet
    from df.enhance import df_features, init_df
    from df.model import ModelParams
    from df.utils import as_complex

from models import dfn3_dir

DF_SR = 48000
CHUNK_SEC = 60.0      # bloque útil por inferencia
CONTEXT_SEC = 5.0     # contexto extra a cada lado para que el modelo "se asiente"
XFADE_SEC = 0.05      # fundido entre bloques
GROUP_SEC = 300.0     # audio máximo por llamada a la GPU
GPU_MIN_SEC = 90.0    # por debajo, la espera de ZeroGPU (cola, asignación) cuesta más que hacerlo en CPU
GPU_COOLDOWN_SEC = 900.0  # tras un fallo de GPU (p. ej. cuota) no se reintenta en este tiempo

ON_ZEROGPU = bool(spaces is not None and spaces.config.Config.zero_gpu)
_gpu_blocked_until = 0.0


def _force_df_device(device: str) -> None:
    """DeepFilterNet decide el dispositivo con df.utils.get_device(); lo fijamos nosotros."""
    dev = torch.device(device)
    for name, mod in list(sys.modules.items()):
        if (name == "df" or name.startswith("df.")) and hasattr(mod, "get_device"):
            mod.get_device = lambda dev=dev: dev


def _disable_df_git() -> None:
    """DeepFilterNet consulta git al iniciar su log; sin git instalado (Windows) eso rompe el arranque."""
    import df.logger
    import df.utils

    for mod in (df.utils, df.logger):
        mod.get_commit_hash = lambda: None
        mod.get_branch_name = lambda: None


print("[MODELO] Cargando DeepFilterNet3…")
_disable_df_git()
_force_df_device("cpu")
_model_cpu, _df_state, _ = init_df(model_base_dir=str(dfn3_dir()), log_level="WARNING", log_file=None)
_model_cpu.eval()
_nb_df = getattr(_model_cpu, "nb_df", getattr(_model_cpu, "df_bins", ModelParams().nb_df))
assert _df_state.sr() == DF_SR

# En ZeroGPU el modelo se coloca en "cuda" a nivel de módulo (emulado fuera de @spaces.GPU).
_model_gpu = None
if ON_ZEROGPU or torch.cuda.is_available():
    try:
        _model_gpu = copy.deepcopy(_model_cpu).to("cuda").eval()
    except Exception as e:  # sin CUDA real ni emulada: solo CPU
        print(f"[MODELO] GPU no disponible: {e}")
print(f"[MODELO] Listo · ZeroGPU={ON_ZEROGPU} · GPU={'sí' if _model_gpu is not None else 'no'}")


@torch.no_grad()
def _enhance_block(model, audio: np.ndarray, atten_db: float | None, device: str) -> np.ndarray:
    """Equivalente a df.enhance.enhance() pero con dispositivo explícito. audio: (C, T) float32 a 48 kHz."""
    x = torch.from_numpy(np.ascontiguousarray(audio, dtype=np.float32))
    orig_len = x.shape[-1]
    n_fft, hop = _df_state.fft_size(), _df_state.hop_size()
    x = F.pad(x, (0, n_fft))  # compensa el retardo del STFT
    spec, erb_feat, spec_feat = df_features(x, _df_state, _nb_df, device=device)
    enhanced = as_complex(model(spec.clone(), erb_feat, spec_feat)[0].cpu().squeeze(1))
    if atten_db is not None and atten_db > 0:
        lim = 10 ** (-atten_db / 20)
        enhanced = as_complex(spec.squeeze(1).cpu()) * lim + enhanced * (1 - lim)
    out = _df_state.synthesis(enhanced.numpy())
    d = n_fft - hop
    return np.asarray(out[:, d : orig_len + d], dtype=np.float32)


def _plan_blocks(n: int) -> list[tuple[int, int, int, int]]:
    """Bloques (ini_ctx, fin_ctx, ini_util, fin_util) en muestras."""
    chunk, ctx = int(CHUNK_SEC * DF_SR), int(CONTEXT_SEC * DF_SR)
    blocks, s = [], 0
    while s < n:
        e = min(n, s + chunk)
        if n - e < chunk // 4:  # no dejar un último bloque diminuto
            e = n
        blocks.append((max(0, s - ctx), min(n, e + ctx), s, e))
        s = e
    return blocks


def _run_blocks(model, pieces: list[np.ndarray], atten_db, device) -> list[np.ndarray]:
    return [_enhance_block(model, p, atten_db, device) for p in pieces]


def _gpu_duration(pieces, atten_db) -> int:
    seconds = sum(p.shape[-1] for p in pieces) / DF_SR
    return int(min(120, 15 + seconds * 0.15))


def _gpu_decorator(fn):
    return spaces.GPU(duration=_gpu_duration)(fn) if spaces is not None else fn


@_gpu_decorator
def _run_blocks_gpu(pieces: list[np.ndarray], atten_db) -> list[np.ndarray]:
    return _run_blocks(_model_gpu, pieces, atten_db, "cuda")


def enhance_48k(audio: np.ndarray, atten_db: float | None, use_gpu: bool = True, progress=None):
    """
    Reduce el ruido de `audio` (C, T) float32 a 48 kHz.
    atten_db: reducción máxima en dB (None = sin límite).
    Devuelve (audio_procesado, info) donde info indica dispositivo y avisos.
    """
    n = audio.shape[-1]
    blocks = _plan_blocks(n)
    xf = int(XFADE_SEC * DF_SR)
    out = np.zeros_like(audio, dtype=np.float32)
    weight = np.zeros(n, dtype=np.float32)

    # Agrupa bloques para que cada llamada a la GPU sea corta.
    per_group = max(1, int(GROUP_SEC // CHUNK_SEC))
    groups = [blocks[i : i + per_group] for i in range(0, len(blocks), per_group)]

    global _gpu_blocked_until
    gpu_ok = (use_gpu and _model_gpu is not None and time.time() >= _gpu_blocked_until
              and (not ON_ZEROGPU or n / DF_SR >= GPU_MIN_SEC))
    devices, notes = set(), []
    for gi, group in enumerate(groups):
        pieces = [audio[:, a:b] for a, b, _, _ in group]
        results = None
        if gpu_ok:
            try:
                results = _run_blocks_gpu(pieces, atten_db)
                devices.add("GPU")
            except Exception as e:
                gpu_ok = False  # si falla una vez (p. ej. cuota), el resto va por CPU
                _gpu_blocked_until = time.time() + GPU_COOLDOWN_SEC
                print(f"[GPU] No disponible, se sigue en CPU: {str(e).splitlines()[0][:200]}")
                notes.append("La tarjeta gráfica no estaba libre: se ha hecho con el procesador (mismo resultado)")
        if results is None:
            results = _run_blocks(_model_cpu, pieces, atten_db, "cpu")
            devices.add("CPU")

        for (a, b, s, e), y in zip(group, results):
            lo, hi = max(0, s - xf), min(n, e + xf)
            w = np.ones(hi - lo, dtype=np.float32)
            if s > 0:  # rampa de entrada centrada en s
                w[: 2 * xf] = np.linspace(0, 1, 2 * xf, dtype=np.float32)[: min(2 * xf, hi - lo)]
            if e < n:  # rampa de salida centrada en e
                w[-2 * xf :] = np.minimum(w[-2 * xf :], np.linspace(1, 0, 2 * xf, dtype=np.float32))
            out[:, lo:hi] += y[:, lo - a : hi - a] * w
            weight[lo:hi] += w
        if progress is not None:
            progress((gi + 1) / len(groups))

    out /= np.maximum(weight, 1e-6)
    return out, {"dispositivo": "+".join(sorted(devices)) or "—", "avisos": notes}


def gpu_available() -> bool:
    return _model_gpu is not None