backend/app/features/audio_preprocessor.py

"""Audio loading, resampling, and normalisation utilities.

Decode chain (in order):
    1. ``soundfile`` (libsndfile)        — fast path for WAV / FLAC / AIFF
    2. ``ffmpeg`` via imageio-ffmpeg     — handles MP3 / OGG / WebM / M4A / etc.

The ffmpeg fallback is critical on Windows because libsndfile ships without
MP3/OGG/WebM codecs and torchaudio's default backend uses libsndfile too.
Browser MediaRecorder produces WebM/Opus, so without this fallback the
microphone-recording path is completely broken.
"""
from __future__ import annotations

import io
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Tuple, Union

import numpy as np
import soundfile as sf
import torch

from app.logging_setup import get_logger

logger = get_logger(__name__)


# --------------------------------------------------------------------------- #
# ffmpeg discovery — prefer the bundled binary, fall back to system PATH.
# --------------------------------------------------------------------------- #

_FFMPEG_EXE: Optional[str] = None


def _resolve_ffmpeg() -> Optional[str]:
    global _FFMPEG_EXE
    if _FFMPEG_EXE is not None:
        return _FFMPEG_EXE
    # Bundled ffmpeg first — works without any system install.
    try:
        import imageio_ffmpeg

        exe = imageio_ffmpeg.get_ffmpeg_exe()
        if exe and Path(exe).exists():
            _FFMPEG_EXE = exe
            logger.info("Using bundled ffmpeg: %s", exe)
            return exe
    except Exception as exc:  # noqa: BLE001
        logger.debug("imageio_ffmpeg unavailable: %s", exc)

    # System ffmpeg as last resort.
    sys_exe = shutil.which("ffmpeg")
    if sys_exe:
        _FFMPEG_EXE = sys_exe
        logger.info("Using system ffmpeg: %s", sys_exe)
        return sys_exe

    logger.warning(
        "No ffmpeg available — only PCM WAV / FLAC / AIFF will decode. "
        "Install imageio-ffmpeg or ffmpeg to enable MP3 / OGG / WebM / M4A."
    )
    return None


# --------------------------------------------------------------------------- #
# Public dataclass
# --------------------------------------------------------------------------- #

@dataclass
class LoadedAudio:
    waveform: torch.Tensor          # [1, T] float32 in ~[-1, 1]
    sample_rate: int                # always == target_sr after preprocessing
    original_sample_rate: int
    original_channels: int
    duration_seconds: float
    file_format: str
    file_size_bytes: int
    decoder: str                    # 'soundfile' | 'ffmpeg' | 'unknown'


# --------------------------------------------------------------------------- #
# Decoders
# --------------------------------------------------------------------------- #

def _decode_with_soundfile(data: Union[bytes, str, Path]) -> Tuple[np.ndarray, int, str]:
    """Try the fast soundfile path. Raises if the format isn't supported."""
    if isinstance(data, (bytes, bytearray)):
        bio = io.BytesIO(data)
        wav, sr = sf.read(bio, dtype="float32", always_2d=True)
        return wav, sr, "wav"
    path = Path(data)
    suffix = path.suffix.lstrip(".").lower() or "unknown"
    wav, sr = sf.read(str(path), dtype="float32", always_2d=True)
    return wav, sr, suffix


def _decode_with_ffmpeg(
    data: Union[bytes, str, Path],
    target_sr: int = 16000,
) -> Tuple[np.ndarray, int, str]:
    """Decode anything ffmpeg understands → mono float32 PCM at target_sr."""
    ffmpeg = _resolve_ffmpeg()
    if ffmpeg is None:
        raise RuntimeError(
            "Audio format not supported by libsndfile and no ffmpeg available. "
            "Install imageio-ffmpeg (pip install imageio-ffmpeg) or ffmpeg."
        )

    is_bytes = isinstance(data, (bytes, bytearray))
    fmt_hint = "compressed"
    if not is_bytes:
        fmt_hint = Path(data).suffix.lstrip(".").lower() or "compressed"

    cmd = [
        ffmpeg,
        "-hide_banner",
        "-loglevel", "error",
        "-nostdin",
        "-i", "pipe:0" if is_bytes else str(data),
        "-vn",                     # discard any video stream
        "-ac", "1",                # mono
        "-ar", str(target_sr),     # resample
        "-f", "wav",
        "-acodec", "pcm_s16le",
        "pipe:1",
    ]

    try:
        proc = subprocess.run(
            cmd,
            input=bytes(data) if is_bytes else None,
            capture_output=True,
            check=False,
            timeout=60,
        )
    except subprocess.TimeoutExpired as exc:
        raise RuntimeError("ffmpeg decode timed out after 60s") from exc

    if proc.returncode != 0:
        msg = proc.stderr.decode(errors="replace").strip() or "unknown ffmpeg error"
        raise RuntimeError(f"ffmpeg decode failed: {msg}")

    bio = io.BytesIO(proc.stdout)
    wav, sr = sf.read(bio, dtype="float32", always_2d=True)
    return wav, sr, fmt_hint


def _read_bytes(
    data: Union[bytes, str, Path],
    target_sr: int = 16000,
) -> Tuple[np.ndarray, int, str, str]:
    """Returns (waveform_np[T,C], sr, format_hint, decoder_used)."""
    # 1) Fast path: soundfile
    try:
        wav, sr, fmt = _decode_with_soundfile(data)
        return wav, sr, fmt, "soundfile"
    except Exception as sf_exc:
        logger.debug("soundfile decode failed: %s; trying ffmpeg.", sf_exc)

    # 2) Fallback: ffmpeg → wav → soundfile
    wav, sr, fmt = _decode_with_ffmpeg(data, target_sr=target_sr)
    return wav, sr, fmt, "ffmpeg"


# --------------------------------------------------------------------------- #
# Public API
# --------------------------------------------------------------------------- #

def preprocess_audio(
    source: Union[bytes, str, Path],
    target_sr: int = 16000,
    max_seconds: float = 30.0,
) -> LoadedAudio:
    """Load → decode → resample → mono → peak-normalize → clip to ``max_seconds``.

    Returns a LoadedAudio with a [1, T] float32 tensor. Raises RuntimeError
    with a clear message on decode failure.
    """
    wav_np, sr, fmt, decoder = _read_bytes(source, target_sr=target_sr)
    file_size = (
        len(source)
        if isinstance(source, (bytes, bytearray))
        else Path(source).stat().st_size
    )
    orig_channels = int(wav_np.shape[1]) if wav_np.ndim == 2 else 1

    # numpy → torch [C, T]
    if wav_np.ndim == 1:
        tensor = torch.from_numpy(wav_np).unsqueeze(0)
    else:
        tensor = torch.from_numpy(wav_np.T)
    tensor = tensor.float()

    # Stereo → mono
    if tensor.shape[0] > 1:
        tensor = tensor.mean(dim=0, keepdim=True)

    # Resample if needed (ffmpeg path is already at target_sr, but be safe).
    if sr != target_sr:
        import torchaudio  # lazy import; not all paths need it

        tensor = torchaudio.functional.resample(tensor, sr, target_sr)

    # Trim to max length
    max_samples = int(max_seconds * target_sr)
    if tensor.shape[1] > max_samples:
        tensor = tensor[:, :max_samples]

    # Peak-normalise (avoid div-by-zero on silence)
    peak = tensor.abs().max()
    if peak > 1e-6:
        tensor = tensor / (peak + 1e-8)

    duration = float(tensor.shape[1]) / target_sr

    return LoadedAudio(
        waveform=tensor.contiguous(),
        sample_rate=target_sr,
        original_sample_rate=sr,
        original_channels=orig_channels,
        duration_seconds=duration,
        file_format=fmt,
        file_size_bytes=int(file_size),
        decoder=decoder,
    )


def waveform_summary(waveform: torch.Tensor, max_points: int = 512) -> list[float]:
    """Down-sample the waveform to ``max_points`` envelope values for the UI."""
    x = waveform.squeeze().detach().cpu().numpy().astype(np.float32)
    if x.size == 0:
        return []
    if x.size <= max_points:
        return x.tolist()
    bucket_size = x.size // max_points
    trimmed = x[: bucket_size * max_points].reshape(max_points, bucket_size)
    env = np.abs(trimmed).max(axis=1)
    return env.tolist()