from __future__ import annotations

import logging
import numpy as np
import soundfile as sf
from pathlib import Path
from typing import Any, Dict
from scipy import signal

from reaper_mcp.core.mcp_core import mcp
from reaper_mcp.core.util import _generate_timestamped_filename

logger = logging.getLogger(__name__)


def apply_fade(
    file_path: str,
    fade_in: float = 0.0,
    fade_out: float = 0.0
) -> Dict[str, Any]:
    """Apply fade in and/or fade out to an audio file.
    
    Args:
        file_path: Path to the input audio file
        fade_in: Fade in duration in seconds (default 0.0)
        fade_out: Fade out duration in seconds (default 0.0)
    
    Returns:
        Dictionary with output file path and fade parameters
    """
    try:
        if fade_in < 0 or fade_out < 0:
            return {"error": "Fade times must be non-negative"}
        
        # Read input file
        audio, sample_rate = sf.read(file_path)
        duration = len(audio) / sample_rate
        
        if fade_in + fade_out > duration:
            return {"error": "Combined fade times exceed audio duration"}
        
        # Apply fade in
        if fade_in > 0:
            fade_in_samples = int(fade_in * sample_rate)
            fade_in_curve = np.linspace(0, 1, fade_in_samples)
            if audio.ndim == 1:
                audio[:fade_in_samples] *= fade_in_curve
            else:
                audio[:fade_in_samples] *= fade_in_curve[:, np.newaxis]
        
        # Apply fade out
        if fade_out > 0:
            fade_out_samples = int(fade_out * sample_rate)
            fade_out_curve = np.linspace(1, 0, fade_out_samples)
            if audio.ndim == 1:
                audio[-fade_out_samples:] *= fade_out_curve
            else:
                audio[-fade_out_samples:] *= fade_out_curve[:, np.newaxis]
        
        # Save to file
        output_path = _generate_timestamped_filename("processed_audio", "faded", "wav")
        sf.write(str(output_path), audio, sample_rate)
        
        logger.info(f"Applied fade to {file_path}: {output_path}")
        return {
            "file_path": str(output_path),
            "input_file": file_path,
            "fade_in": fade_in,
            "fade_out": fade_out,
            "duration": duration,
            "sample_rate": sample_rate
        }
    except Exception as e:
        logger.error(f"Failed to apply fade: {e}", exc_info=True)
        return {"error": f"Failed to apply fade: {e}"}


def reverse_audio(file_path: str) -> Dict[str, Any]:
    """Reverse an audio file (play backwards).
    
    Args:
        file_path: Path to the input audio file
    
    Returns:
        Dictionary with output file path
    """
    try:
        # Read input file
        audio, sample_rate = sf.read(file_path)
        
        # Reverse audio
        if audio.ndim == 1:
            reversed_audio = audio[::-1]
        else:
            reversed_audio = audio[::-1, :]
        
        # Save to file
        output_path = _generate_timestamped_filename("processed_audio", "reversed", "wav")
        sf.write(str(output_path), reversed_audio, sample_rate)
        
        logger.info(f"Reversed audio {file_path}: {output_path}")
        return {
            "file_path": str(output_path),
            "input_file": file_path,
            "duration": len(audio) / sample_rate,
            "sample_rate": sample_rate
        }
    except Exception as e:
        logger.error(f"Failed to reverse audio: {e}", exc_info=True)
        return {"error": f"Failed to reverse audio: {e}"}


def time_stretch_audio(
    file_path: str,
    rate: float,
    preserve_pitch: bool = False
) -> Dict[str, Any]:
    """Time stretch audio (change speed/duration).
    
    Args:
        file_path: Path to the input audio file
        rate: Time stretch rate (0.5 = half speed, 2.0 = double speed)
        preserve_pitch: If True, maintains original pitch (default False, requires librosa)
    
    Returns:
        Dictionary with output file path and stretch parameters
    """
    try:
        if rate <= 0 or rate > 10:
            return {"error": "Rate must be between 0 and 10"}
        
        # Read input file
        audio, sample_rate = sf.read(file_path)
        
        # For simple time stretching without pitch preservation,
        # we can resample. For pitch preservation, we'd need librosa
        # which isn't in our dependencies yet. For now, implement
        # simple resampling (changes pitch).
        
        if preserve_pitch:
            return {"error": "Pitch-preserving time stretch requires librosa library (not yet installed)"}
        
        # Simple time stretch by resampling (changes pitch)
        from scipy.signal import resample
        
        new_length = int(len(audio) / rate)
        if audio.ndim == 1:
            stretched_audio = resample(audio, new_length)
        else:
            stretched_audio = np.zeros((new_length, audio.shape[1]))
            for channel in range(audio.shape[1]):
                stretched_audio[:, channel] = resample(audio[:, channel], new_length)
        
        # Save to file
        output_path = _generate_timestamped_filename("processed_audio", "stretched", "wav")
        sf.write(str(output_path), stretched_audio, sample_rate)
        
        logger.info(f"Time stretched {file_path} by {rate}x: {output_path}")
        return {
            "file_path": str(output_path),
            "input_file": file_path,
            "rate": rate,
            "preserve_pitch": preserve_pitch,
            "original_duration": len(audio) / sample_rate,
            "new_duration": len(stretched_audio) / sample_rate,
            "sample_rate": sample_rate
        }
    except Exception as e:
        logger.error(f"Failed to time stretch audio: {e}", exc_info=True)
        return {"error": f"Failed to time stretch audio: {e}"}


def pitch_shift_audio(
    file_path: str,
    semitones: float
) -> Dict[str, Any]:
    """Shift the pitch of audio by semitones.
    
    Args:
        file_path: Path to the input audio file
        semitones: Number of semitones to shift (positive = up, negative = down)
    
    Returns:
        Dictionary with output file path and shift parameters
    """
    try:
        if semitones < -24 or semitones > 24:
            return {"error": "Semitones must be between -24 and 24"}
        
        # Read input file
        audio, sample_rate = sf.read(file_path)
        
        # Simple pitch shift using resampling
        # This changes the duration as well (not ideal, but works without librosa)
        pitch_ratio = 2 ** (semitones / 12)
        
        from scipy.signal import resample
        
        new_length = int(len(audio) / pitch_ratio)
        if audio.ndim == 1:
            shifted_audio = resample(audio, new_length)
        else:
            shifted_audio = np.zeros((new_length, audio.shape[1]))
            for channel in range(audio.shape[1]):
                shifted_audio[:, channel] = resample(audio[:, channel], new_length)
        
        # Save to file
        output_path = _generate_timestamped_filename("processed_audio", "pitched", "wav")
        sf.write(str(output_path), shifted_audio, sample_rate)
        
        logger.info(f"Pitch shifted {file_path} by {semitones} semitones: {output_path}")
        return {
            "file_path": str(output_path),
            "input_file": file_path,
            "semitones": semitones,
            "original_duration": len(audio) / sample_rate,
            "new_duration": len(shifted_audio) / sample_rate,
            "sample_rate": sample_rate,
            "note": "Duration changed due to simple pitch shifting algorithm"
        }
    except Exception as e:
        logger.error(f"Failed to pitch shift audio: {e}", exc_info=True)
        return {"error": f"Failed to pitch shift audio: {e}"}


def normalize_audio(
    file_path: str,
    target_db: float = -3.0
) -> Dict[str, Any]:
    """Normalize audio to a target peak level in dB.
    
    Args:
        file_path: Path to the input audio file
        target_db: Target peak level in dB (default -3.0)
    
    Returns:
        Dictionary with output file path and normalization info
    """
    try:
        if target_db > 0:
            return {"error": "Target dB must be 0 or negative to avoid clipping"}
        if target_db < -60:
            return {"error": "Target dB must be greater than -60"}
        
        # Read input file
        audio, sample_rate = sf.read(file_path)
        
        # Calculate current peak level
        current_peak = np.max(np.abs(audio))
        
        if current_peak == 0:
            return {"error": "Audio is silent, cannot normalize"}
        
        # Calculate target amplitude from dB
        target_amplitude = 10 ** (target_db / 20)
        
        # Calculate gain needed
        gain = target_amplitude / current_peak
        
        # Apply gain
        normalized_audio = audio * gain
        
        # Save to file
        output_path = _generate_timestamped_filename("processed_audio", "normalized", "wav")
        sf.write(str(output_path), normalized_audio, sample_rate)
        
        logger.info(f"Normalized {file_path} to {target_db} dB: {output_path}")
        return {
            "file_path": str(output_path),
            "input_file": file_path,
            "target_db": target_db,
            "gain_applied": gain,
            "gain_db": 20 * np.log10(gain),
            "sample_rate": sample_rate
        }
    except Exception as e:
        logger.error(f"Failed to normalize audio: {e}", exc_info=True)
        return {"error": f"Failed to normalize audio: {e}"}


def apply_eq(
    file_path: str,
    low_gain: float = 0.0,
    mid_gain: float = 0.0,
    high_gain: float = 0.0
) -> Dict[str, Any]:
    """Apply 3-band EQ to audio file.
    
    Args:
        file_path: Path to the input audio file
        low_gain: Low frequency gain in dB (-12 to +12, default 0.0)
        mid_gain: Mid frequency gain in dB (-12 to +12, default 0.0)
        high_gain: High frequency gain in dB (-12 to +12, default 0.0)
    
    Returns:
        Dictionary with output file path and EQ parameters
    """
    try:
        if not (-12 <= low_gain <= 12 and -12 <= mid_gain <= 12 and -12 <= high_gain <= 12):
            return {"error": "All gain values must be between -12 and +12 dB"}
        
        # Read input file
        audio, sample_rate = sf.read(file_path)
        
        # Define frequency bands (low: <250Hz, mid: 250Hz-4kHz, high: >4kHz)
        low_cutoff = 250
        high_cutoff = 4000
        
        # Convert gains from dB to linear
        low_multiplier = 10 ** (low_gain / 20)
        mid_multiplier = 10 ** (mid_gain / 20)
        high_multiplier = 10 ** (high_gain / 20)
        
        # Design filters
        nyquist = sample_rate / 2
        
        # Low band (lowpass)
        b_low, a_low = signal.butter(4, low_cutoff / nyquist, btype='low')
        
        # Mid band (bandpass)
        b_mid, a_mid = signal.butter(4, [low_cutoff / nyquist, high_cutoff / nyquist], btype='band')
        
        # High band (highpass)
        b_high, a_high = signal.butter(4, high_cutoff / nyquist, btype='high')
        
        # Apply filters to each channel
        if audio.ndim == 1:
            low_band = signal.filtfilt(b_low, a_low, audio) * low_multiplier
            mid_band = signal.filtfilt(b_mid, a_mid, audio) * mid_multiplier
            high_band = signal.filtfilt(b_high, a_high, audio) * high_multiplier
            eq_audio = low_band + mid_band + high_band
        else:
            eq_audio = np.zeros_like(audio)
            for channel in range(audio.shape[1]):
                low_band = signal.filtfilt(b_low, a_low, audio[:, channel]) * low_multiplier
                mid_band = signal.filtfilt(b_mid, a_mid, audio[:, channel]) * mid_multiplier
                high_band = signal.filtfilt(b_high, a_high, audio[:, channel]) * high_multiplier
                eq_audio[:, channel] = low_band + mid_band + high_band
        
        # Prevent clipping
        max_val = np.max(np.abs(eq_audio))
        if max_val > 1.0:
            eq_audio = eq_audio / max_val
        
        # Save to file
        output_path = _generate_timestamped_filename("processed_audio", "eq", "wav")
        sf.write(str(output_path), eq_audio, sample_rate)
        
        logger.info(f"Applied EQ to {file_path}: {output_path}")
        return {
            "file_path": str(output_path),
            "input_file": file_path,
            "low_gain": low_gain,
            "mid_gain": mid_gain,
            "high_gain": high_gain,
            "sample_rate": sample_rate
        }
    except Exception as e:
        logger.error(f"Failed to apply EQ: {e}", exc_info=True)
        return {"error": f"Failed to apply EQ: {e}"}


def apply_compression(
    file_path: str,
    threshold: float = -20.0,
    ratio: float = 4.0,
    attack: float = 0.005,
    release: float = 0.1
) -> Dict[str, Any]:
    """Apply dynamic range compression to audio file.
    
    Args:
        file_path: Path to the input audio file
        threshold: Threshold in dB below which no compression occurs (default -20.0)
        ratio: Compression ratio (default 4.0, meaning 4:1)
        attack: Attack time in seconds (default 0.005)
        release: Release time in seconds (default 0.1)
    
    Returns:
        Dictionary with output file path and compression parameters
    """
    try:
        if threshold > 0:
            return {"error": "Threshold must be 0 or negative"}
        if ratio < 1.0 or ratio > 20.0:
            return {"error": "Ratio must be between 1.0 and 20.0"}
        if attack <= 0 or release <= 0:
            return {"error": "Attack and release must be positive"}
        
        # Read input file
        audio, sample_rate = sf.read(file_path)
        
        # Simple compression implementation
        # Convert threshold from dB to linear
        threshold_linear = 10 ** (threshold / 20)
        
        # Calculate envelope (RMS-based)
        if audio.ndim == 1:
            audio_mono = audio
        else:
            audio_mono = np.mean(audio, axis=1)
        
        # Calculate RMS envelope with sliding window
        window_size = int(0.01 * sample_rate)  # 10ms window
        rms = np.sqrt(np.convolve(audio_mono ** 2, np.ones(window_size) / window_size, mode='same'))
        
        # Calculate gain reduction
        gain = np.ones_like(rms)
        above_threshold = rms > threshold_linear
        
        # Apply compression ratio to signals above threshold
        gain[above_threshold] = (threshold_linear / rms[above_threshold]) ** ((ratio - 1) / ratio)
        
        # Smooth gain changes with attack and release
        attack_samples = int(attack * sample_rate)
        release_samples = int(release * sample_rate)
        
        smoothed_gain = np.copy(gain)
        for i in range(1, len(gain)):
            if gain[i] < smoothed_gain[i-1]:
                # Attack (gain reduction)
                alpha = 1.0 - np.exp(-1.0 / attack_samples)
            else:
                # Release (gain recovery)
                alpha = 1.0 - np.exp(-1.0 / release_samples)
            
            smoothed_gain[i] = alpha * gain[i] + (1 - alpha) * smoothed_gain[i-1]
        
        # Apply gain to audio
        if audio.ndim == 1:
            compressed_audio = audio * smoothed_gain
        else:
            compressed_audio = audio * smoothed_gain[:, np.newaxis]
        
        # Save to file
        output_path = _generate_timestamped_filename("processed_audio", "compressed", "wav")
        sf.write(str(output_path), compressed_audio, sample_rate)
        
        logger.info(f"Applied compression to {file_path}: {output_path}")
        return {
            "file_path": str(output_path),
            "input_file": file_path,
            "threshold": threshold,
            "ratio": ratio,
            "attack": attack,
            "release": release,
            "sample_rate": sample_rate
        }
    except Exception as e:
        logger.error(f"Failed to apply compression: {e}", exc_info=True)
        return {"error": f"Failed to apply compression: {e}"}


# MCP Tool Wrappers
# These wrapper functions expose the audio processing utilities as MCP tools

@mcp.tool()
def audio_apply_fade(
    file_path: str,
    fade_in: float = 0.0,
    fade_out: float = 0.0
) -> Dict[str, Any]:
    """Apply fade in and/or fade out to an audio file.
    
    Args:
        file_path: Path to the input audio file
        fade_in: Fade in duration in seconds (default 0.0)
        fade_out: Fade out duration in seconds (default 0.0)
    
    Returns:
        Dictionary with output file path and fade parameters
    """
    return apply_fade(file_path, fade_in, fade_out)


@mcp.tool()
def audio_reverse(file_path: str) -> Dict[str, Any]:
    """Reverse an audio file (play backwards).
    
    Args:
        file_path: Path to the input audio file
    
    Returns:
        Dictionary with output file path
    """
    return reverse_audio(file_path)


@mcp.tool()
def audio_time_stretch(
    file_path: str,
    rate: float,
    preserve_pitch: bool = False
) -> Dict[str, Any]:
    """Time stretch audio (change speed/duration).
    
    Args:
        file_path: Path to the input audio file
        rate: Time stretch rate (0.5 = half speed, 2.0 = double speed)
        preserve_pitch: If True, maintains original pitch (default False, requires librosa)
    
    Returns:
        Dictionary with output file path and stretch parameters
    """
    return time_stretch_audio(file_path, rate, preserve_pitch)


@mcp.tool()
def audio_pitch_shift(
    file_path: str,
    semitones: float
) -> Dict[str, Any]:
    """Shift the pitch of audio by semitones.
    
    Args:
        file_path: Path to the input audio file
        semitones: Number of semitones to shift (positive = up, negative = down)
    
    Returns:
        Dictionary with output file path and shift parameters
    """
    return pitch_shift_audio(file_path, semitones)


@mcp.tool()
def audio_normalize(
    file_path: str,
    target_db: float = -3.0
) -> Dict[str, Any]:
    """Normalize audio to a target peak level in dB.
    
    Args:
        file_path: Path to the input audio file
        target_db: Target peak level in dB (default -3.0)
    
    Returns:
        Dictionary with output file path and normalization info
    """
    return normalize_audio(file_path, target_db)


@mcp.tool()
def audio_apply_eq(
    file_path: str,
    low_gain: float = 0.0,
    mid_gain: float = 0.0,
    high_gain: float = 0.0
) -> Dict[str, Any]:
    """Apply 3-band EQ to audio file.
    
    Args:
        file_path: Path to the input audio file
        low_gain: Low frequency gain in dB (-12 to +12, default 0.0)
        mid_gain: Mid frequency gain in dB (-12 to +12, default 0.0)
        high_gain: High frequency gain in dB (-12 to +12, default 0.0)
    
    Returns:
        Dictionary with output file path and EQ parameters
    """
    return apply_eq(file_path, low_gain, mid_gain, high_gain)


@mcp.tool()
def audio_apply_compression(
    file_path: str,
    threshold: float = -20.0,
    ratio: float = 4.0,
    attack: float = 0.005,
    release: float = 0.1
) -> Dict[str, Any]:
    """Apply dynamic range compression to audio file.
    
    Args:
        file_path: Path to the input audio file
        threshold: Threshold in dB below which no compression occurs (default -20.0)
        ratio: Compression ratio (default 4.0, meaning 4:1)
        attack: Attack time in seconds (default 0.005)
        release: Release time in seconds (default 0.1)
    
    Returns:
        Dictionary with output file path and compression parameters
    """
    return apply_compression(file_path, threshold, ratio, attack, release)
