"""Audio cleanup for camera mics, using speechbrain.

Pipeline per audio frame (16 kHz mono int16, ~100 ms chunks from the mic):
  int16 -> float32 [-1, 1] -> speechbrain enhancement -> int16

Two implementations, selected by availability and the ENHANCE_MODE config:
  * "model"  - speechbrain SpectralMaskEnhancement (metricGAN-plus). Real
               ML-based noise suppression. Runs on CPU; ~20-40 ms per 100 ms
               chunk on a modern core, so latency stays well under a frame
               budget if workers are spread across cores.
  * "gain"   - lightweight DSP fallback (soft-clip + high-pass) used when the
               model can't be loaded (ENHANCE_MODE=off / auto + no model).

The enhancer is created once per worker process and reused across frames.
"""

from __future__ import annotations

import logging
from typing import Optional

import numpy as np

log = logging.getLogger("audio.cleanup")

SAMPLE_RATE = 16000


class AudioCleaner:
    """Stateful audio cleaner. Call .process(int16 mono) -> int16 mono."""

    def __init__(self, mode: str = "auto",
                 model_source: str = "speechbrain/metricgan-plus-voicebank",
                 sample_rate: int = SAMPLE_RATE,
                 chunk_s: float = 1.0,
                 hop_s: float = 0.2):
        self.mode = mode
        self.sample_rate = sample_rate
        self.chunk = max(1, int(chunk_s * sample_rate))
        self.hop = max(1, int(hop_s * sample_rate))
        self._ctx = np.zeros(0, dtype=np.float32)
        self._enhancer = None
        self._torch = None
        self.backend = "off"

        if mode == "force":
            self._load_model(model_source)
            if self.backend == "off":
                raise RuntimeError(
                    f"ENHANCE_MODE=force but failed to load model "
                    f"'{model_source}'")
        elif mode == "auto":
            self._load_model(model_source)
        # mode == "off": passthrough + basic gain shaping

    def _load_model(self, source: str) -> None:
        try:
            import torch
            from speechbrain.inference.enhancement import (
                SpectralMaskEnhancement,
            )
        except Exception as e:  # noqa: BLE001
            log.info("speechbrain unavailable (%s); using fallback", e)
            self._enhancer = None
            return
        try:
            log.info("loading speechbrain enhancement model %s ...", source)
            self._torch = torch
            model = SpectralMaskEnhancement.from_hparams(
                source=source,
                savedir="/root/.cache/speechbrain-enhancement",
            )
            model.mods.enhance_model.eval()
            self._enhancer = model
            self.backend = "model"
            log.info("speechbrain enhancement ready (16 kHz in/out)")
        except Exception as e:  # noqa: BLE001
            log.warning("failed to load enhancement model %s: %s "
                        "(using fallback)", source, e)
            self._enhancer = None

    # -- DSP fallback ------------------------------------------------------

    @staticmethod
    def _simple_clean(x: np.ndarray) -> np.ndarray:
        """Lightweight cleanup: high-pass + soft knee + gentle gain.

        Good enough to remove DC offset/low rumble and taming clipping
        without any ML model.
        """
        # first-order high-pass at ~80 Hz
        if len(x) > 1:
            alpha = 0.998
            y = np.empty_like(x)
            y[0] = x[0]
            for i in range(1, len(x)):
                y[i] = alpha * (y[i - 1] + x[i] - x[i - 1])
            x = y
        # soft-clip to avoid harsh distortion
        x = np.tanh(x * 0.9) / np.tanh(0.9)
        # gentle makeup gain
        x = x * 1.1
        return np.clip(x, -1.0, 1.0)

    # -- public API ---------------------------------------------------------

    def process(self, pcm: np.ndarray) -> np.ndarray:
        """pcm: int16 mono numpy array. Returns int16 mono, same length.

        Accumulates a `chunk_s` window and runs the model on that window,
        emitting only the latest hop (the input length). Until a full window
        is buffered, the DSP fallback is used so the stream is not silent.
        """
        n = len(pcm)
        if n == 0:
            return pcm
        x = pcm.astype(np.float32) / 32768.0

        if self._enhancer is None:
            y = self._simple_clean(x)
            return (np.clip(y, -1.0, 1.0) * 32767.0).astype(np.int16)

        self._ctx = np.concatenate([self._ctx, x])
        if len(self._ctx) < self.chunk:
            y = self._simple_clean(x)
            return (np.clip(y, -1.0, 1.0) * 32767.0).astype(np.int16)
        self._ctx = self._ctx[-self.chunk:]

        torch = self._torch
        try:
            w = torch.from_numpy(self._ctx.copy()).unsqueeze(0).to(
                self._enhancer.device)
            # speechbrain >= 1.1 (e.g. metricGAN-plus) requires `lengths`
            # on enhance_batch. The value is RELATIVE (multiplied by the
            # number of frames internally), so 1.0 = "the whole sequence".
            # Older versions accept the call without it (TypeError fallback).
            lengths = torch.tensor([1.0]).to(self._enhancer.device)
            try:
                with torch.no_grad():
                    out = self._enhancer.enhance_batch(w, lengths=lengths)
            except TypeError:
                with torch.no_grad():
                    out = self._enhancer.enhance_batch(w)
            enhanced = out.squeeze(0).cpu().numpy()
        except Exception as e:  # noqa: BLE001
            log.warning("enhance_batch failed (%s); using fallback", e)
            y = self._simple_clean(x)
            return (np.clip(y, -1.0, 1.0) * 32767.0).astype(np.int16)

        hop_out = enhanced[-n:] if len(enhanced) >= n else np.pad(
            enhanced, (0, n - len(enhanced)))
        hop_out = hop_out[:n]
        return (np.clip(hop_out, -1.0, 1.0) * 32767.0).astype(np.int16)