"""Audio processing utilities for SalTTS."""

from typing import Optional, Tuple

import librosa
import numpy as np
import soundfile as sf
import torch


def load_audio(
    path: str,
    sample_rate: int = 24000,
    normalize: bool = True,
    trim_silence: bool = True,
) -> Tuple[np.ndarray, int]:
    """Load audio file.

    Args:
        path: Path to audio file
        sample_rate: Target sample rate
        normalize: Whether to normalize audio
        trim_silence: Whether to trim silence

    Returns:
        Tuple of (audio waveform, sample rate)
    """
    audio, sr = librosa.load(path, sr=sample_rate, mono=True)

    if trim_silence:
        audio, _ = librosa.effects.trim(audio, top_db=20)

    if normalize:
        audio = audio / np.max(np.abs(audio) + 1e-8)

    return audio, sr


def save_audio(path: str, audio: np.ndarray, sample_rate: int = 24000) -> None:
    """Save audio to file.

    Args:
        path: Output path
        audio: Audio waveform
        sample_rate: Sample rate
    """
    sf.write(path, audio, sample_rate)


def compute_mel_spectrogram(
    audio: np.ndarray,
    sample_rate: int = 24000,
    n_fft: int = 1024,
    hop_length: int = 256,
    win_length: int = 1024,
    n_mels: int = 80,
    fmin: float = 0.0,
    fmax: float = 8000.0,
) -> np.ndarray:
    """Compute mel spectrogram.

    Args:
        audio: Audio waveform
        sample_rate: Sample rate
        n_fft: FFT size
        hop_length: Hop length
        win_length: Window length
        n_mels: Number of mel bins
        fmin: Minimum frequency
        fmax: Maximum frequency

    Returns:
        Mel spectrogram
    """
    mel_spec = librosa.feature.melspectrogram(
        y=audio,
        sr=sample_rate,
        n_fft=n_fft,
        hop_length=hop_length,
        win_length=win_length,
        n_mels=n_mels,
        fmin=fmin,
        fmax=fmax,
    )
    return librosa.power_to_db(mel_spec, ref=np.max)


def normalize_audio(audio: np.ndarray, target_db: float = -20.0) -> np.ndarray:
    """Normalize audio to target dB level.

    Args:
        audio: Audio waveform
        target_db: Target dB level

    Returns:
        Normalized audio
    """
    rms = np.sqrt(np.mean(audio**2))
    target_rms = 10 ** (target_db / 20)
    return audio * (target_rms / (rms + 1e-8))


def apply_fade(audio: np.ndarray, fade_duration_ms: float = 2.0, sample_rate: int = 24000) -> np.ndarray:
    """Apply fade in/out to audio.

    Args:
        audio: Audio waveform
        fade_duration_ms: Fade duration in milliseconds
        sample_rate: Sample rate

    Returns:
        Audio with fades applied
    """
    fade_samples = int(fade_duration_ms * sample_rate / 1000)
    fade_in = np.linspace(0, 1, fade_samples)
    fade_out = np.linspace(1, 0, fade_samples)

    audio[:fade_samples] *= fade_in
    audio[-fade_samples:] *= fade_out

    return audio
