"""
Audio Processing Utilities

Audio processing and conversion utilities using pydub and librosa

License: BSD 3-Clause
"""

from pathlib import Path
from typing import Union, Optional, Tuple
import numpy as np

from pydub import AudioSegment
from cocina.printer import Printer
printer = Printer()

try:
    import librosa
    LIBROSA_AVAILABLE = True
except ImportError:
    LIBROSA_AVAILABLE = False

try:
    from scipy import signal
    SCIPY_AVAILABLE = True
except ImportError:
    SCIPY_AVAILABLE = False


def pydub_audio_to_numpy(audio: Union[str, AudioSegment]) -> Tuple[np.ndarray, int]:
    """
    Convert audio file or AudioSegment to numpy array using pydub.
    
    Args:
        audio: Either a file path (str) or a pydub AudioSegment object
        
    Returns:
        Tuple of (audio_np_array, sample_rate)
        - audio_np_array: Normalized numpy array with values in [-1.0, 1.0]
        - sample_rate: Sample rate in Hz
    """
    if isinstance(audio, str):
        audio = AudioSegment.from_file(audio)
    
    # Get raw audio data
    raw_data = audio.raw_data
    sample_rate = audio.frame_rate
    channels = audio.channels
    sample_width = audio.sample_width
    
    # Convert to numpy array based on sample width
    if sample_width == 1:
        # 8-bit audio (unsigned) - special case
        audio_array = np.frombuffer(raw_data, dtype=np.uint8).astype(np.float32)
        audio_array = (audio_array - 128.0) / 128.0  # Convert unsigned to signed, then normalize
    elif sample_width == 2:
        # 16-bit audio (signed)
        audio_array = np.frombuffer(raw_data, dtype=np.int16).astype(np.float32)
        audio_array = audio_array / (2**(8 * sample_width - 1))  # 2^15 = 32768
    elif sample_width == 3:
        # 24-bit audio (signed) - pad to 32-bit
        raw_data_padded = b''.join([b'\x00' + raw_data[i:i+3] for i in range(0, len(raw_data), 3)])
        audio_array = np.frombuffer(raw_data_padded, dtype=np.int32).astype(np.float32)
        audio_array = audio_array / (2**(8 * sample_width - 1))  # 2^23 = 8388608
    elif sample_width == 4:
        # 32-bit audio (signed)
        audio_array = np.frombuffer(raw_data, dtype=np.int32).astype(np.float32)
        audio_array = audio_array / (2**(8 * sample_width - 1))  # 2^31 = 2147483648
    else:
        raise ValueError(f"Unsupported sample width: {sample_width} bytes")
    
    # Handle multi-channel audio
    if channels > 1:
        audio_array = audio_array.reshape(-1, channels)
        # Convert to mono by averaging channels
        audio_array = np.mean(audio_array, axis=1)
    
    printer.message(f"Converted audio: shape={audio_array.shape}, sample_rate={sample_rate}Hz, "
                f"channels={channels}, sample_width={sample_width}bytes")
    
    return audio_array, sample_rate


def librosa_audio_to_numpy(audio: Union[str, AudioSegment], 
                          sr: Optional[int] = None) -> Tuple[np.ndarray, int]:
    """
    Convert audio file or AudioSegment to numpy array using librosa.
    
    Args:
        audio: Either a file path (str) or a pydub AudioSegment object
        sr: Target sample rate (if None, uses original sample rate)
        
    Returns:
        Tuple of (audio_np_array, sample_rate)
        - audio_np_array: Normalized numpy array with values in [-1.0, 1.0]
        - sample_rate: Sample rate in Hz
        
    Raises:
        ImportError: If librosa is not available
    """
    if not LIBROSA_AVAILABLE:
        raise ImportError("librosa is not available. Install it using: pixi add librosa")
    
    if isinstance(audio, AudioSegment):
        # Convert AudioSegment to temporary file for librosa to read
        import tempfile
        import os
        
        with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file:
            temp_path = temp_file.name
            audio.export(temp_path, format="wav")
        
        try:
            audio_array, sample_rate = librosa.load(temp_path, sr=sr, mono=True)
        finally:
            os.unlink(temp_path)
    else:
        # Load directly from file path
        audio_array, sample_rate = librosa.load(audio, sr=sr, mono=True)
    
    printer.message(f"Loaded audio with librosa: shape={audio_array.shape}, sample_rate={sample_rate}Hz")
    
    return audio_array, sample_rate


def amplitude_to_db(S: np.ndarray, 
                   ref: float = 1.0,
                   amin: float = 1e-10,
                   top_db: Optional[float] = 80.0,
                   dest: Optional[str] = None,
                   create_parents: bool = True,
                   overwrite: bool = True) -> Union[np.ndarray, str]:
    """
    Convert amplitude spectrogram to dB-scale spectrogram using librosa.
    
    Args:
        S: Input amplitude spectrogram
        ref: Reference value for dB calculation
        amin: Minimum amplitude threshold
        top_db: Maximum dB range (None for no limit)
        dest: Optional destination file path to save result
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing destination file
        
    Returns:
        dB-scale spectrogram array if dest is None, otherwise path to saved file
        
    Raises:
        ImportError: If librosa is not available
    """
    if not LIBROSA_AVAILABLE:
        raise ImportError("librosa is not available. Install it using: pixi add librosa")
    
    S_db = librosa.amplitude_to_db(S, ref=ref, amin=amin, top_db=top_db)
    
    if dest is None:
        return S_db
    
    # Save to file
    dest_path = Path(dest)
    if create_parents:
        dest_path.parent.mkdir(parents=True, exist_ok=True)
    
    if dest_path.exists() and not overwrite:
        raise FileExistsError(f"Destination file {dest} already exists and overwrite=False")
    
    np.save(dest_path, S_db)
    printer.message(f"Saved amplitude_to_db result to: {dest_path}")
    
    return str(dest_path)


def spectrogram(y: np.ndarray,
               sr: int,
               n_fft: int = 2048,
               hop_length: Optional[int] = None,
               win_length: Optional[int] = None,
               window: str = 'hann',
               dest: Optional[str] = None,
               create_parents: bool = True,
               overwrite: bool = True) -> Union[np.ndarray, str]:
    """
    Compute short-time Fourier transform (STFT) spectrogram using librosa.
    
    Args:
        y: Audio time series
        sr: Sample rate
        n_fft: FFT window size
        hop_length: Number of samples between frames (default: n_fft//4)
        win_length: Window length (default: n_fft)
        window: Window function type
        dest: Optional destination file path to save result
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing destination file
        
    Returns:
        Complex-valued STFT matrix if dest is None, otherwise path to saved file
        
    Raises:
        ImportError: If librosa is not available
    """
    if not LIBROSA_AVAILABLE:
        raise ImportError("librosa is not available. Install it using: pixi add librosa")
    
    if hop_length is None:
        hop_length = n_fft // 4
    
    S = librosa.stft(y, n_fft=n_fft, hop_length=hop_length, 
                     win_length=win_length, window=window)
    
    if dest is None:
        return S
    
    # Save to file
    dest_path = Path(dest)
    if create_parents:
        dest_path.parent.mkdir(parents=True, exist_ok=True)
    
    if dest_path.exists() and not overwrite:
        raise FileExistsError(f"Destination file {dest} already exists and overwrite=False")
    
    np.save(dest_path, S)
    printer.message(f"Saved spectrogram result to: {dest_path}")
    
    return str(dest_path)


def mel_spectrogram(y: np.ndarray,
                   sr: int,
                   n_fft: int = 2048,
                   hop_length: Optional[int] = None,
                   win_length: Optional[int] = None,
                   window: str = 'hann',
                   n_mels: int = 128,
                   fmin: float = 0.0,
                   fmax: Optional[float] = None,
                   dest: Optional[str] = None,
                   create_parents: bool = True,
                   overwrite: bool = True) -> Union[np.ndarray, str]:
    """
    Compute mel-scaled spectrogram using librosa.
    
    Args:
        y: Audio time series
        sr: Sample rate
        n_fft: FFT window size
        hop_length: Number of samples between frames (default: n_fft//4)
        win_length: Window length (default: n_fft)
        window: Window function type
        n_mels: Number of mel bands
        fmin: Minimum frequency
        fmax: Maximum frequency (default: sr/2)
        dest: Optional destination file path to save result
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing destination file
        
    Returns:
        Mel spectrogram array if dest is None, otherwise path to saved file
        
    Raises:
        ImportError: If librosa is not available
    """
    if not LIBROSA_AVAILABLE:
        raise ImportError("librosa is not available. Install it using: pixi add librosa")
    
    if hop_length is None:
        hop_length = n_fft // 4
    
    if fmax is None:
        fmax = sr / 2.0
    
    S = librosa.feature.melspectrogram(y=y, sr=sr, n_fft=n_fft, 
                                      hop_length=hop_length, win_length=win_length,
                                      window=window, n_mels=n_mels, 
                                      fmin=fmin, fmax=fmax)
    
    if dest is None:
        return S
    
    # Save to file
    dest_path = Path(dest)
    if create_parents:
        dest_path.parent.mkdir(parents=True, exist_ok=True)
    
    if dest_path.exists() and not overwrite:
        raise FileExistsError(f"Destination file {dest} already exists and overwrite=False")
    
    np.save(dest_path, S)
    printer.message(f"Saved mel spectrogram result to: {dest_path}")
    
    return str(dest_path)


def pydub_spectrogram(audio: Union[str, AudioSegment],
                     nperseg: int = 2048,
                     noverlap: Optional[int] = None,
                     nfft: Optional[int] = None,
                     dest: Optional[str] = None,
                     create_parents: bool = True,
                     overwrite: bool = True) -> Union[Tuple[np.ndarray, np.ndarray, np.ndarray], str]:
    """
    Compute spectrogram using pydub and scipy.signal.
    
    Args:
        audio: Either a file path (str) or a pydub AudioSegment object
        nperseg: Length of each segment for FFT
        noverlap: Number of points to overlap between segments (default: nperseg//8)
        nfft: Length of the FFT used (default: nperseg)
        dest: Optional destination file path to save result
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing destination file
        
    Returns:
        Tuple of (frequencies, times, Sxx) if dest is None, otherwise path to saved file
        - frequencies: Array of sample frequencies
        - times: Array of segment times
        - Sxx: Spectrogram of audio
        
    Raises:
        ImportError: If scipy is not available
    """
    if not SCIPY_AVAILABLE:
        raise ImportError("scipy is not available. Install it using: pixi add scipy")
    
    # Convert to numpy array
    audio_array, sample_rate = pydub_audio_to_numpy(audio)
    
    if noverlap is None:
        noverlap = nperseg // 8
    
    if nfft is None:
        nfft = nperseg
    
    # Compute spectrogram using scipy
    frequencies, times, Sxx = signal.spectrogram(
        audio_array, 
        fs=sample_rate,
        nperseg=nperseg,
        noverlap=noverlap,
        nfft=nfft
    )
    
    if dest is None:
        return frequencies, times, Sxx
    
    # Save to file (save all three arrays)
    dest_path = Path(dest)
    if create_parents:
        dest_path.parent.mkdir(parents=True, exist_ok=True)
    
    if dest_path.exists() and not overwrite:
        raise FileExistsError(f"Destination file {dest} already exists and overwrite=False")
    
    # Save as npz file with all three arrays
    np.savez(dest_path, frequencies=frequencies, times=times, Sxx=Sxx)
    printer.message(f"Saved pydub spectrogram result to: {dest_path}")
    
    return str(dest_path)


def pydub_amplitude_to_db(S: np.ndarray,
                         ref: float = 1.0,
                         amin: float = 1e-10,
                         top_db: Optional[float] = 80.0,
                         dest: Optional[str] = None,
                         create_parents: bool = True,
                         overwrite: bool = True) -> Union[np.ndarray, str]:
    """
    Convert amplitude spectrogram to dB-scale using numpy (pydub-compatible version).
    
    Args:
        S: Input amplitude spectrogram
        ref: Reference value for dB calculation
        amin: Minimum amplitude threshold
        top_db: Maximum dB range (None for no limit)
        dest: Optional destination file path to save result
        create_parents: Create parent directories if needed
        overwrite: Overwrite existing destination file
        
    Returns:
        dB-scale spectrogram array if dest is None, otherwise path to saved file
    """
    # Ensure S is positive and above minimum threshold
    S = np.maximum(amin, np.abs(S))
    
    # Convert to dB scale: 20 * log10(S / ref)
    S_db = 20.0 * np.log10(S / ref)
    
    # Apply top_db limit if specified
    if top_db is not None:
        S_db = np.maximum(S_db, S_db.max() - top_db)
    
    if dest is None:
        return S_db
    
    # Save to file
    dest_path = Path(dest)
    if create_parents:
        dest_path.parent.mkdir(parents=True, exist_ok=True)
    
    if dest_path.exists() and not overwrite:
        raise FileExistsError(f"Destination file {dest} already exists and overwrite=False")
    
    np.save(dest_path, S_db)
    printer.message(f"Saved pydub amplitude_to_db result to: {dest_path}")
    
    return str(dest_path)