"""Audio preprocessing pipeline for XLSR-SLS deepfake detection.

Transforms input audio into 16 kHz mono float32 with fixed or windowed frames.
Robustly supports microphone recordings (WebM/Opus), raw numpy arrays, and standard audio files.
"""
from __future__ import annotations

import math
import os
from pathlib import Path
import subprocess
import tempfile
import uuid
import numpy as np
import soundfile as sf
from scipy.signal import resample_poly

SAMPLE_RATE = 16_000
WINDOW = 64_600            # samples required by the XLSR-SLS model (~4.04 s)
CLIP_SECONDS = WINDOW / SAMPLE_RATE


def load(audio_input: str | Path | tuple | dict | np.ndarray) -> tuple[np.ndarray, int]:
    """Read an audio file or microphone buffer -> (float32 samples, sample_rate)."""
    # 1. Handle Gradio (sample_rate, numpy_array) tuple
    if isinstance(audio_input, tuple) and len(audio_input) == 2:
        sr, data = audio_input
        data = np.asarray(data, dtype=np.float32)
        # If integer PCM, normalize to [-1.0, 1.0]
        if data.dtype in (np.int16, np.int32):
            max_val = np.iinfo(data.dtype).max
            data = data / max_val
        elif np.max(np.abs(data)) > 1.0:
            data = data / 32768.0
        return data, int(sr)

    # 2. Handle dictionary format
    if isinstance(audio_input, dict) and "data" in audio_input and "sample_rate" in audio_input:
        sr = audio_input["sample_rate"]
        data = np.asarray(audio_input["data"], dtype=np.float32)
        if np.max(np.abs(data)) > 1.0:
            data = data / 32768.0
        return data, int(sr)

    # 3. Handle file paths (WAV, MP3, WebM, OGG, M4A, etc.)
    p = Path(audio_input)
    if not p.is_file():
        raise FileNotFoundError(f"Audio file not found: {audio_input}")

    # Try 3.1: soundfile directly
    try:
        data, sr = sf.read(str(p), dtype="float32", always_2d=True)
        return data, sr
    except Exception:
        pass

    # Try 3.2: librosa
    try:
        import librosa
        data, sr = librosa.load(str(p), sr=None, mono=False)
        if data.ndim == 1:
            data = data[:, None]
        else:
            data = data.T  # (samples, channels)
        return data.astype(np.float32), int(sr)
    except Exception:
        pass

    # Try 3.3: ffmpeg system CLI (converts WebM/Opus or any container to temporary WAV)
    temp_wav = Path(tempfile.gettempdir()) / f"conv_{uuid.uuid4().hex[:8]}.wav"
    try:
        cmd = [
            "ffmpeg", "-y", "-i", str(p),
            "-vn", "-ac", "1", "-ar", "16000",
            "-f", "wav", str(temp_wav)
        ]
        res = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, timeout=10)
        if res.returncode == 0 and temp_wav.is_file():
            data, sr = sf.read(str(temp_wav), dtype="float32", always_2d=True)
            return data, sr
    except Exception:
        pass
    finally:
        if temp_wav.is_file():
            try:
                temp_wav.unlink()
            except OSError:
                pass

    raise RuntimeError(
        f"Could not decode audio from {p.name}. Ensure audio format is valid (WAV, MP3, WebM, FLAC, OGG)."
    )


def to_mono(data: np.ndarray) -> np.ndarray:
    """Downmix channels to mono."""
    if data.ndim == 1:
        return data
    return data.mean(axis=1)


def resample(audio: np.ndarray, sr_in: int, sr_out: int = SAMPLE_RATE) -> np.ndarray:
    if sr_in == sr_out:
        return audio
    g = math.gcd(int(sr_in), int(sr_out))
    return resample_poly(audio, sr_out // g, sr_in // g).astype(np.float32)


def fit_length(audio: np.ndarray, length: int = WINDOW, crop: str = "start") -> np.ndarray:
    """Force audio to exactly `length` samples. Tile if shorter; crop if longer."""
    n = len(audio)
    if n == 0:
        raise ValueError("Audio clip is empty (0 samples). Please speak into the microphone or upload a valid recording.")
    if n < 1600:  # less than 0.1s
        raise ValueError("Audio recording is too short (< 0.1s). Please speak for at least 1-2 seconds.")
    if n == length:
        return audio
    if n < length:
        return np.tile(audio, math.ceil(length / n))[:length]
    if crop == "start":
        s = 0
    elif crop == "center":
        s = (n - length) // 2
    elif crop == "loudest":
        power = np.concatenate([[0.0], np.cumsum(audio.astype(np.float64) ** 2)])
        s = int(np.argmax(power[length:] - power[:-length]))
    else:
        s = 0
    return audio[s:s + length]


def peak_normalize(audio: np.ndarray, target: float = 0.97) -> np.ndarray:
    peak = float(np.max(np.abs(audio)))
    return audio * (target / peak) if peak > 0 else audio


def preprocess_array(audio: np.ndarray, sr: int, crop: str = "start", normalize: bool = False) -> np.ndarray:
    audio = to_mono(np.asarray(audio, dtype=np.float32))
    audio = resample(audio, sr, SAMPLE_RATE)
    audio = fit_length(audio, WINDOW, crop)
    if normalize:
        audio = peak_normalize(audio)
    return np.ascontiguousarray(audio, dtype=np.float32)


def preprocess_file(path: str | Path, crop: str = "start", normalize: bool = False) -> np.ndarray:
    data, sr = load(path)
    return preprocess_array(data, sr, crop=crop, normalize=normalize)