app/core/audio_processor.py

"""Audio Processing Module.

Handles all audio-related operations:
- Base64 decoding
- MP3/WAV conversion
- Resampling to 16kHz
- Duration validation
"""

import base64
import io
import numpy as np
from typing import Tuple
from pydub import AudioSegment
import librosa

from app.config import settings


class AudioProcessingError(Exception):
    """Base exception for audio processing errors."""
    pass


class InvalidBase64Error(AudioProcessingError):
    """Raised when Base64 decoding fails."""
    pass


class InvalidAudioFormatError(AudioProcessingError):
    """Raised when audio format is invalid."""
    pass


class AudioTooShortError(AudioProcessingError):
    """Raised when audio is too short."""
    pass


class AudioTooLongError(AudioProcessingError):
    """Raised when audio is too long."""
    pass


class AudioProcessor:
    """Handles audio processing operations."""
    
    def __init__(self, target_sr: int = None):
        """Initialize audio processor.
        
        Args:
            target_sr: Target sample rate (default from settings: 16000)
        """
        self.target_sr = target_sr or settings.TARGET_SAMPLE_RATE
    
    def decode_base64(self, b64_string: str) -> bytes:
        """Decode Base64 string to bytes.
        
        Args:
            b64_string: Base64 encoded audio string
            
        Returns:
            Raw audio bytes
            
        Raises:
            InvalidBase64Error: If decoding fails
        """
        try:
            return base64.b64decode(b64_string)
        except Exception as e:
            raise InvalidBase64Error(f"Failed to decode Base64: {str(e)}")
    
    def load_audio_from_bytes(self, audio_bytes: bytes) -> AudioSegment:
        """Load audio from bytes (supports MP3, WAV).
        
        Args:
            audio_bytes: Raw audio file bytes
            
        Returns:
            AudioSegment object
            
        Raises:
            InvalidAudioFormatError: If format is not supported
        """
        try:
            # Try MP3 first
            audio = AudioSegment.from_mp3(io.BytesIO(audio_bytes))
            return audio
        except Exception:
            pass
        
        try:
            # Try WAV
            audio = AudioSegment.from_wav(io.BytesIO(audio_bytes))
            return audio
        except Exception:
            pass
        
        try:
            # Try generic format detection (ffmpeg handles M4A, MOV, etc.)
            audio = AudioSegment.from_file(io.BytesIO(audio_bytes))
            return audio
        except Exception as e:
            raise InvalidAudioFormatError(
                f"Could not decode audio. Supported formats: MP3, WAV, M4A, AAC. Error: {str(e)}"
            )
    
    def validate_duration(self, audio: AudioSegment) -> float:
        """Validate audio duration is within limits.
        
        Args:
            audio: AudioSegment object
            
        Returns:
            Duration in seconds
            
        Raises:
            AudioTooShortError: If duration < MIN_DURATION
            AudioTooLongError: If duration > MAX_DURATION
        """
        duration = len(audio) / 1000.0  # Convert ms to seconds
        
        if duration < settings.MIN_DURATION:
            raise AudioTooShortError(
                f"Audio too short: {duration:.2f}s (minimum: {settings.MIN_DURATION}s)"
            )
        
        if duration > settings.MAX_DURATION:
            raise AudioTooLongError(
                f"Audio too long: {duration:.2f}s (maximum: {settings.MAX_DURATION}s)"
            )
        
        return duration
    
    def convert_to_wav_buffer(self, audio: AudioSegment) -> io.BytesIO:
        """Convert AudioSegment to WAV buffer.
        
        Args:
            audio: AudioSegment object
            
        Returns:
            BytesIO buffer containing WAV data
        """
        wav_buffer = io.BytesIO()
        audio.export(wav_buffer, format="wav")
        wav_buffer.seek(0)
        return wav_buffer
    
    def resample_to_numpy(self, wav_buffer: io.BytesIO) -> np.ndarray:
        """Load WAV and resample to target sample rate.
        
        Args:
            wav_buffer: BytesIO buffer with WAV data
            
        Returns:
            Numpy array of audio samples (mono, resampled)
        """
        waveform, _ = librosa.load(
            wav_buffer, 
            sr=self.target_sr,
            mono=True
        )
        # 1. Trim Silence (DISABLED: Removing silence makes real speech sound unnaturally continuous/AI-like)
        # waveform, _ = librosa.effects.trim(waveform, top_db=30)
        
        # 2. Normalize Volume (Re-enabled: Necessary for consistent chunked inference)
        waveform = librosa.util.normalize(waveform)

        # 3. Dithering / Noise Injection (DISABLED)
        # Reason: Dithering added noise that helped Human audio pass, but also helped High-Quality AI pass.
        # We now rely on the Heuristic Override to catch compressed human audio, so we don't need to fake the noise.
        # noise_amp = 0.005 * np.max(np.abs(waveform))
        # waveform = waveform + noise_amp * np.random.normal(size=len(waveform))
        
        return waveform
    
    def process(self, b64_audio: str) -> Tuple[np.ndarray, float]:
        """Full audio processing pipeline (Base64 input).
        
        Decodes Base64 -> Loads audio -> Validates -> Converts -> Resamples
        """
        # Step 1: Decode Base64
        audio_bytes = self.decode_base64(b64_audio)
        return self.process_bytes(audio_bytes)

    def process_bytes(self, audio_bytes: bytes) -> Tuple[np.ndarray, float]:
        """Full audio processing pipeline (Bytes input).
        
        Loads audio -> Validates -> Converts -> Resamples
        
        Args:
            audio_bytes: Raw audio bytes
            
        Returns:
            Tuple of (waveform numpy array, duration in seconds)
        """
        # Step 2: Load audio
        audio_segment = self.load_audio_from_bytes(audio_bytes)
        
        # Step 3: Validate duration
        duration = self.validate_duration(audio_segment)
        
        # Step 4: Convert to WAV
        wav_buffer = self.convert_to_wav_buffer(audio_segment)
        
        # Step 5: Resample to numpy array
        waveform = self.resample_to_numpy(wav_buffer)
        
        return waveform, duration


# Create singleton instance
audio_processor = AudioProcessor()