#!/usr/bin/env python3
"""
Audio Processor - Handles audio loading, preprocessing and caching
===================================================================
"""

import os
import hashlib
import logging
import numpy as np
import soundfile as sf
from pathlib import Path
from typing import Tuple, Optional
import subprocess
import tempfile

logger = logging.getLogger(__name__)


class AudioProcessor:
    """Handles audio file processing and caching"""
    
    def __init__(self, cache_dir: str = "/tmp/audio_cache", target_sr: int = 16000):
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.target_sr = target_sr
        
    def load_audio(self, audio_path: str, use_cache: bool = True) -> Tuple[np.ndarray, int]:
        """Load audio file and convert to proper format"""
        
        audio_path = Path(audio_path)
        
        if not audio_path.exists():
            raise FileNotFoundError(f"Audio file not found: {audio_path}")
        
        # Check cache
        if use_cache:
            cached_audio = self._get_cached_audio(audio_path)
            if cached_audio is not None:
                logger.info(f"Loaded audio from cache: {audio_path.name}")
                return cached_audio, self.target_sr
        
        # Load audio based on extension
        audio_data = None
        sample_rate = None
        
        if audio_path.suffix.lower() in ['.wav', '.flac']:
            # Direct load for supported formats
            audio_data, sample_rate = sf.read(str(audio_path))
            
        elif audio_path.suffix.lower() in ['.mp3', '.m4a', '.ogg', '.webm']:
            # Convert using ffmpeg for other formats
            audio_data, sample_rate = self._convert_with_ffmpeg(audio_path)
            
        else:
            raise ValueError(f"Unsupported audio format: {audio_path.suffix}")
        
        # Convert to mono if stereo
        if len(audio_data.shape) > 1:
            audio_data = np.mean(audio_data, axis=1)
        
        # Resample if needed
        if sample_rate != self.target_sr:
            audio_data = self._resample_audio(audio_data, sample_rate, self.target_sr)
            sample_rate = self.target_sr
        
        # Normalize audio
        audio_data = self._normalize_audio(audio_data)
        
        # Cache the processed audio
        if use_cache:
            self._cache_audio(audio_path, audio_data)
        
        return audio_data, sample_rate
    
    def _convert_with_ffmpeg(self, audio_path: Path) -> Tuple[np.ndarray, int]:
        """Convert audio using ffmpeg"""
        
        logger.info(f"Converting {audio_path.suffix} file with ffmpeg...")
        
        with tempfile.NamedTemporaryFile(suffix='.wav') as tmp_wav:
            # Convert to WAV using ffmpeg
            cmd = [
                'ffmpeg', '-i', str(audio_path),
                '-ar', str(self.target_sr),
                '-ac', '1',  # Mono
                '-f', 'wav',
                '-y',  # Overwrite
                tmp_wav.name
            ]
            
            try:
                result = subprocess.run(
                    cmd, 
                    capture_output=True, 
                    text=True, 
                    timeout=30
                )
                
                if result.returncode != 0:
                    logger.error(f"FFmpeg error: {result.stderr}")
                    raise RuntimeError(f"FFmpeg conversion failed: {result.stderr}")
                
                # Load the converted WAV
                audio_data, sample_rate = sf.read(tmp_wav.name)
                return audio_data, sample_rate
                
            except subprocess.TimeoutExpired:
                logger.error("FFmpeg conversion timed out")
                raise
            except FileNotFoundError:
                logger.error("FFmpeg not found. Please install ffmpeg")
                raise RuntimeError("FFmpeg not found. Install with: apt-get install ffmpeg")
    
    def _resample_audio(self, audio: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
        """Resample audio to target sample rate"""
        
        if orig_sr == target_sr:
            return audio
        
        # Simple resampling using numpy
        duration = len(audio) / orig_sr
        n_samples = int(duration * target_sr)
        
        # Use linear interpolation for resampling
        x_old = np.linspace(0, duration, len(audio))
        x_new = np.linspace(0, duration, n_samples)
        
        resampled = np.interp(x_new, x_old, audio)
        
        return resampled
    
    def _normalize_audio(self, audio: np.ndarray) -> np.ndarray:
        """Normalize audio to [-1, 1] range"""
        
        # Avoid division by zero
        max_val = np.max(np.abs(audio))
        if max_val > 0:
            audio = audio / max_val * 0.95  # Leave some headroom
        
        return audio.astype(np.float32)
    
    def _get_cache_key(self, audio_path: Path) -> str:
        """Generate cache key for audio file"""
        
        # Use file path and modification time for cache key
        stat = audio_path.stat()
        key_str = f"{audio_path}_{stat.st_mtime}_{stat.st_size}"
        return hashlib.md5(key_str.encode()).hexdigest()
    
    def _get_cached_audio(self, audio_path: Path) -> Optional[np.ndarray]:
        """Retrieve cached audio if available"""
        
        cache_key = self._get_cache_key(audio_path)
        cache_file = self.cache_dir / f"{cache_key}.npy"
        
        if cache_file.exists():
            try:
                return np.load(cache_file)
            except Exception as e:
                logger.warning(f"Failed to load cached audio: {e}")
                cache_file.unlink()  # Remove corrupted cache
        
        return None
    
    def _cache_audio(self, audio_path: Path, audio_data: np.ndarray):
        """Save processed audio to cache"""
        
        cache_key = self._get_cache_key(audio_path)
        cache_file = self.cache_dir / f"{cache_key}.npy"
        
        try:
            np.save(cache_file, audio_data)
            logger.debug(f"Cached audio: {cache_key}")
        except Exception as e:
            logger.warning(f"Failed to cache audio: {e}")
    
    def clear_cache(self):
        """Clear all cached audio files"""
        
        cache_files = list(self.cache_dir.glob("*.npy"))
        for cache_file in cache_files:
            cache_file.unlink()
        
        logger.info(f"Cleared {len(cache_files)} cached audio files")
    
    def preprocess_for_whisper(self, audio_data: np.ndarray) -> np.ndarray:
        """Preprocess audio specifically for Whisper model"""
        
        # Whisper expects float32 audio in [-1, 1] range
        if audio_data.dtype != np.float32:
            audio_data = audio_data.astype(np.float32)
        
        # Ensure proper range
        audio_data = np.clip(audio_data, -1.0, 1.0)
        
        return audio_data