"""
Audio decoder module for Sempress audio compression.

Reconstructs audio from compressed representation.
"""

from __future__ import annotations
from typing import Tuple, Dict
import numpy as np
from .container import unpack_container

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


def decode_audio(compressed_blob: bytes) -> Tuple[np.ndarray, int]:
    """
    Decode compressed audio blob back to waveform.
    
    Args:
        compressed_blob: Compressed audio data
        
    Returns:
        audio: (n_samples,) or (n_samples, n_channels) waveform
        sample_rate: int
    """
    # Unpack container
    payload = unpack_container(compressed_blob)
    
    if payload["domain"] != "audio":
        raise ValueError(f"Expected audio domain, got {payload['domain']}")
    
    # Reconstruct each channel
    n_channels = payload["n_channels"]
    sr = payload["sample_rate"]
    n_samples = payload["n_samples"]
    config = payload["config"]
    
    channels = []
    for ch_idx in range(n_channels):
        ch_data = payload["channel_data"][f"ch{ch_idx}"]
        ch_audio = _decode_audio_channel(ch_data, n_samples, config)
        channels.append(ch_audio)
    
    # Combine channels
    if n_channels == 1:
        audio = channels[0]
    else:
        # Stack channels: (n_samples, n_channels)
        audio = np.stack(channels, axis=1)
    
    return audio, sr


def _decode_audio_channel(ch_data: Dict, n_samples: int, config: Dict) -> np.ndarray:
    """Decode a single audio channel."""
    # Load codebook
    codebook_shape = ch_data["codebook_shape"]
    codebook = np.frombuffer(ch_data["codebook"], dtype=np.float32).reshape(codebook_shape)
    
    # Load indices
    indices = np.frombuffer(ch_data["indices"], dtype=np.uint16)
    
    # Reconstruct features from codebook
    features_recon = codebook[indices]
    
    # Apply residuals if present
    residuals = ch_data.get("residuals", {})
    if residuals and "indices" in residuals:
        residual_indices = np.frombuffer(residuals["indices"], dtype=np.uint32)
        residual_data = np.frombuffer(residuals["data"], dtype=np.float32)
        residual_data = residual_data.reshape(len(residual_indices), -1)
        features_recon[residual_indices] += residual_data
    
    # Convert features back to audio
    feature_metadata = ch_data["feature_metadata"]
    
    if feature_metadata["feature_type"] == "mfcc":
        # Inverse MFCC to waveform (approximate reconstruction)
        audio = _mfcc_to_audio(
            features_recon.T,
            n_samples=n_samples,
            n_fft=config["frame_size"],
            hop_length=config["hop_length"]
        )
    else:
        # Inverse frame to waveform (overlap-add)
        audio = _frames_to_audio(
            features_recon.T,
            hop_length=config["hop_length"],
            n_samples=n_samples
        )
    
    return audio


def _mfcc_to_audio(mfcc: np.ndarray, n_samples: int, n_fft: int, hop_length: int) -> np.ndarray:
    """
    Convert MFCC features back to audio waveform.
    This is an approximate reconstruction using inverse mel-frequency transform.
    """
    if not LIBROSA_AVAILABLE:
        raise RuntimeError("librosa is required for MFCC decoding")
    
    # Inverse MFCC to mel-spectrogram
    mel_spectrogram = librosa.feature.inverse.mfcc_to_mel(mfcc)
    
    # Inverse mel-spectrogram to linear spectrogram (magnitude)
    mel_basis = librosa.filters.mel(sr=22050, n_fft=n_fft)  # Default sr, will be stretched
    mel_basis_inv = np.linalg.pinv(mel_basis)
    magnitude = mel_basis_inv @ mel_spectrogram
    
    # Griffin-Lim algorithm to estimate phase and reconstruct audio
    audio = librosa.griffinlim(
        magnitude,
        n_iter=32,
        hop_length=hop_length,
        n_fft=n_fft
    )
    
    # Trim or pad to correct length
    if len(audio) > n_samples:
        audio = audio[:n_samples]
    elif len(audio) < n_samples:
        audio = np.pad(audio, (0, n_samples - len(audio)))
    
    return audio


def _frames_to_audio(frames: np.ndarray, hop_length: int, n_samples: int) -> np.ndarray:
    """
    Convert frames back to audio using overlap-add.
    
    Args:
        frames: (frame_size, n_frames)
        hop_length: hop length between frames
        n_samples: target number of samples
        
    Returns:
        audio: (n_samples,)
    """
    frame_size, n_frames = frames.shape
    
    # Overlap-add reconstruction
    audio_length = (n_frames - 1) * hop_length + frame_size
    audio = np.zeros(audio_length)
    window = np.hanning(frame_size)
    
    for i in range(n_frames):
        start = i * hop_length
        audio[start:start + frame_size] += frames[:, i] * window
    
    # Trim or pad to correct length
    if len(audio) > n_samples:
        audio = audio[:n_samples]
    elif len(audio) < n_samples:
        audio = np.pad(audio, (0, n_samples - len(audio)))
    
    return audio


def decode_to_audio_file(compressed_blob: bytes, output_path: str) -> Dict:
    """
    Decode compressed audio and save to file.
    
    Args:
        compressed_blob: Compressed audio data
        output_path: Output audio file path (e.g., 'output.wav')
        
    Returns:
        Metadata dictionary with reconstruction info
    """
    try:
        import soundfile as sf
    except ImportError:
        raise RuntimeError("soundfile is required to save audio files. Install with: pip install soundfile")
    
    # Decode audio
    audio, sr = decode_audio(compressed_blob)
    
    # Save to file
    sf.write(output_path, audio, sr)
    
    # Extract metadata
    payload = unpack_container(compressed_blob)
    
    return {
        "sample_rate": sr,
        "n_samples": len(audio) if audio.ndim == 1 else audio.shape[0],
        "n_channels": 1 if audio.ndim == 1 else audio.shape[1],
        "method": payload["config"]["method"],
        "use_spectral": payload["config"]["use_spectral"]
    }
