from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np
import soundfile as sf
from scipy import signal
from scipy.fft import rfft, rfftfreq

from reaper_mcp.core.mcp_core import mcp
from reaper_mcp.core.util import _load_sample_metadata, _save_sample_metadata, is_cache_valid

logger = logging.getLogger(__name__)


def _estimate_tempo(audio: np.ndarray, sr: int) -> float:
    """Estimate tempo (BPM) using autocorrelation method.
    
    This is a basic implementation using autocorrelation to find periodicity.
    More sophisticated methods would use librosa or essentia.
    """
    try:
        # Check for silence/very low energy first
        rms_energy = np.sqrt(np.mean(audio ** 2))
        if rms_energy < 1e-6:  # Threshold for silence
            return 120.0  # Default for silence
        
        # Use envelope for tempo detection
        envelope = np.abs(signal.hilbert(audio))
        
        # Downsample envelope for efficiency
        hop_length = 512
        envelope_ds = envelope[::hop_length]
        
        # Compute autocorrelation
        autocorr = np.correlate(envelope_ds, envelope_ds, mode='full')
        autocorr = autocorr[len(autocorr)//2:]
        
        # Find peaks in autocorrelation
        # BPM range: 60-180 BPM (typical range)
        min_period = int(60 / 180 * sr / hop_length)  # samples for 180 BPM
        max_period = int(60 / 60 * sr / hop_length)   # samples for 60 BPM
        
        if max_period >= len(autocorr):
            max_period = len(autocorr) - 1
        
        if min_period >= max_period:
            return 120.0  # Default fallback
        
        # Find strongest peak in valid range
        search_range = autocorr[min_period:max_period]
        if len(search_range) == 0:
            return 120.0
            
        peak_idx = np.argmax(search_range) + min_period
        
        # Convert to BPM
        tempo = 60 * sr / (peak_idx * hop_length)
        
        # Clamp to reasonable range
        tempo = max(60.0, min(180.0, tempo))
        
        return round(tempo, 1)
    except Exception as e:
        logger.warning(f"Tempo estimation failed: {e}")
        return 120.0  # Default fallback


def _estimate_key(audio: np.ndarray, sr: int) -> str:
    """Estimate musical key using spectral analysis.
    
    This is a simplified implementation. More accurate key detection
    would require librosa or essentia with chromagram analysis.
    """
    try:
        # Use a portion of the audio (middle section, 5 seconds max)
        duration = len(audio) / sr
        start = int(sr * duration * 0.3)
        end = int(start + min(sr * 5, len(audio) - start))
        audio_slice = audio[start:end]
        
        # Compute FFT
        fft_result = rfft(audio_slice)
        freqs = rfftfreq(len(audio_slice), 1/sr)
        magnitudes = np.abs(fft_result)
        
        # Focus on musical frequency range (27.5 Hz - 4186 Hz, A0 to C8)
        valid_freq_mask = (freqs >= 27.5) & (freqs <= 4186)
        freqs = freqs[valid_freq_mask]
        magnitudes = magnitudes[valid_freq_mask]
        
        if len(magnitudes) == 0:
            return "C"
        
        # Check for silence (very low energy)
        max_magnitude = np.max(magnitudes)
        if max_magnitude < 1e-6:  # Threshold for silence
            return "C"
        
        # Find dominant frequency
        peak_idx = np.argmax(magnitudes)
        dominant_freq = freqs[peak_idx]
        
        # Convert frequency to note (simplified chromatic scale)
        # A4 = 440 Hz
        notes = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
        
        # Calculate semitones from A4
        if dominant_freq > 0:
            semitones_from_a4 = 12 * np.log2(dominant_freq / 440.0)
            note_idx = int(round(semitones_from_a4)) % 12
            # Adjust to start from C (A is at index 9 in the notes array)
            note_idx = (note_idx + 9) % 12
            return notes[note_idx]
        
        return "C"
    except Exception as e:
        logger.warning(f"Key estimation failed: {e}")
        return "C"


def _analyze_audio_data(audio: np.ndarray, sr: int, file_path: str) -> Dict[str, Any]:
    """Perform analysis on audio data and return metadata."""
    try:
        # Handle multi-channel audio
        if audio.ndim > 1:
            # Convert to mono for analysis
            audio_mono = np.mean(audio, axis=1)
            channels = audio.shape[1]
        else:
            audio_mono = audio
            channels = 1
        
        duration = len(audio_mono) / sr
        
        # Calculate RMS energy
        rms_energy = float(np.sqrt(np.mean(audio_mono ** 2)))
        
        # Calculate peak amplitude
        peak_amplitude = float(np.max(np.abs(audio_mono)))
        
        # Calculate zero crossing rate (spectral characteristic)
        zero_crossings = np.sum(np.diff(np.sign(audio_mono)) != 0)
        zcr = zero_crossings / len(audio_mono)
        
        # Estimate tempo
        tempo = _estimate_tempo(audio_mono, sr)
        
        # Estimate key
        key = _estimate_key(audio_mono, sr)
        
        # Get file modification time for cache invalidation
        mtime = Path(file_path).stat().st_mtime
        
        metadata = {
            "file_path": str(file_path),
            "duration": round(duration, 3),
            "sample_rate": int(sr),
            "channels": int(channels),
            "rms_energy": round(rms_energy, 6),
            "peak_amplitude": round(peak_amplitude, 6),
            "zero_crossing_rate": round(zcr, 6),
            "estimated_tempo_bpm": tempo,
            "estimated_key": key,
            "mtime": mtime,
        }
        
        return metadata
        
    except Exception as e:
        logger.error(f"Audio analysis failed for {file_path}: {e}")
        raise


def _analyze_audio_file(file_path: str) -> Dict[str, Any]:
    """Internal implementation: Analyze an audio file and extract comprehensive metadata.
    
    Extracts: BPM (tempo), key, duration, sample rate, channels, RMS energy,
    peak amplitude, zero-crossing rate (spectral characteristic).
    
    Args:
        file_path: Full path to the audio file to analyze
        
    Returns:
        Dictionary containing audio metadata and analysis results
        
    Note: This uses scipy/numpy for analysis. For more accurate tempo and key detection,
          consider installing librosa or essentia (currently not included due to complexity).
    """
    try:
        if not Path(file_path).is_file():
            return {"error": f"File not found: {file_path}"}
        
        # Check cache first
        cache = _load_sample_metadata()
        cached_data = cache.get(file_path)
        
        if cached_data and is_cache_valid(file_path, cached_data.get("mtime")):
            logger.info(f"Using cached analysis for {file_path}")
            return {"analysis": cached_data, "cached": True}
        
        # Load audio file
        logger.info(f"Analyzing audio file: {file_path}")
        audio, sr = sf.read(file_path)
        
        # Perform analysis
        metadata = _analyze_audio_data(audio, sr, file_path)
        
        # Update cache
        cache[file_path] = metadata
        _save_sample_metadata(cache)
        
        return {"analysis": metadata, "cached": False}
        
    except Exception as e:
        error_msg = f"Failed to analyze audio file: {e}"
        logger.error(error_msg, exc_info=True)
        return {"error": error_msg}


@mcp.tool()
def analyze_audio_file(file_path: str) -> Dict[str, Any]:
    """Analyze an audio file and extract comprehensive metadata.
    
    Extracts: BPM (tempo), key, duration, sample rate, channels, RMS energy,
    peak amplitude, zero-crossing rate (spectral characteristic).
    
    Args:
        file_path: Full path to the audio file to analyze
        
    Returns:
        Dictionary containing audio metadata and analysis results
        
    Note: This uses scipy/numpy for analysis. For more accurate tempo and key detection,
          consider installing librosa or essentia (currently not included due to complexity).
    """
    return _analyze_audio_file(file_path)


def _batch_analyze_samples(
    directory: str,
    extensions: Optional[List[str]] = None,
    force_reanalyze: bool = False
) -> Dict[str, Any]:
    """Internal implementation: Analyze all audio samples in a directory and create searchable metadata.
    
    Recursively scans the directory for audio files, analyzes them, and caches results
    for fast retrieval. Uses cached results when available unless force_reanalyze is True.
    
    Args:
        directory: Path to directory containing audio samples
        extensions: List of file extensions to analyze (default: ['.wav', '.aiff', '.aif', '.flac', '.mp3', '.ogg'])
        force_reanalyze: If True, reanalyze all files even if cached (default: False)
        
    Returns:
        Dictionary with analysis results, including count of analyzed and cached files
    """
    try:
        dir_path = Path(directory)
        if not dir_path.is_dir():
            return {"error": f"Directory not found: {directory}"}
        
        if not extensions:
            extensions = ['.wav', '.aiff', '.aif', '.flac', '.mp3', '.ogg']
        
        # Normalize extensions to lowercase
        extensions = [ext.lower() if ext.startswith('.') else f'.{ext.lower()}' for ext in extensions]
        
        # Load existing cache
        cache = _load_sample_metadata()
        
        analyzed_count = 0
        cached_count = 0
        error_count = 0
        results = []
        
        # Scan directory recursively
        logger.info(f"Scanning directory: {directory}")
        for file_path_obj in dir_path.rglob('*'):
            if file_path_obj.is_file() and any(file_path_obj.name.lower().endswith(ext) for ext in extensions):
                file_path = str(file_path_obj)
                
                try:
                    # Check cache
                    cached_data = cache.get(file_path)
                    
                    if not force_reanalyze and cached_data and is_cache_valid(file_path, cached_data.get("mtime")):
                        # Use cached data
                        cached_count += 1
                        results.append(cached_data)
                    else:
                        # Analyze file
                        logger.info(f"Analyzing: {file_path}")
                        audio, sr = sf.read(file_path)
                        metadata = _analyze_audio_data(audio, sr, file_path)
                        
                        cache[file_path] = metadata
                        analyzed_count += 1
                        results.append(metadata)
                        
                except Exception as e:
                    logger.warning(f"Failed to analyze {file_path}: {e}")
                    error_count += 1
        
        # Save updated cache
        _save_sample_metadata(cache)
        
        logger.info(f"Batch analysis complete: {analyzed_count} analyzed, {cached_count} cached, {error_count} errors")
        
        return {
            "directory": directory,
            "total_files": len(results),
            "analyzed": analyzed_count,
            "cached": cached_count,
            "errors": error_count,
            "results": results,
        }
        
    except Exception as e:
        error_msg = f"Failed to batch analyze samples: {e}"
        logger.error(error_msg, exc_info=True)
        return {"error": error_msg}


@mcp.tool()
def batch_analyze_samples(
    directory: str,
    extensions: Optional[List[str]] = None,
    force_reanalyze: bool = False
) -> Dict[str, Any]:
    """Analyze all audio samples in a directory and create searchable metadata.
    
    Recursively scans the directory for audio files, analyzes them, and caches results
    for fast retrieval. Uses cached results when available unless force_reanalyze is True.
    
    Args:
        directory: Path to directory containing audio samples
        extensions: List of file extensions to analyze (default: ['.wav', '.aiff', '.aif', '.flac', '.mp3', '.ogg'])
        force_reanalyze: If True, reanalyze all files even if cached (default: False)
        
    Returns:
        Dictionary with analysis results, including count of analyzed and cached files
    """
    return _batch_analyze_samples(directory, extensions, force_reanalyze)
