"""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):
        self.mode = mode
        self.sample_rate = sample_rate
        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.

        If the model changes the length, we trim/pad back to the original
        length so downstream chunking stays aligned.
        """
        if self._enhancer is None:
            if len(pcm) == 0:
                return pcm
            x = pcm.astype(np.float32) / 32768.0
            y = self._simple_clean(x)
            return (y * 32767.0).astype(np.float32).astype(np.int16)

        x = pcm.astype(np.float32) / 32768.0
        torch = self._torch
        try:
            w = torch.from_numpy(x).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)
            out = 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 (y * 32767.0).astype(np.float32).astype(np.int16)

        # align length
        n = len(x)
        if len(out) >= n:
            out = out[:n]
        else:
            out = np.pad(out, (0, n - len(out)))
        return (out * 32767.0).astype(np.float32).astype(np.int16)