audio_processor.py

"""
Audio Volume Normalizer - Audio Processing Module

This module handles audio processing operations including:
- Loading audio files (MP3 and WAV)
- Detecting loudness levels
- Normalizing audio volume
"""

import logging
import os
from typing import Dict, List, Tuple, Optional, Callable
import numpy as np
from pydub import AudioSegment

logger = logging.getLogger(__name__)


class AudioProcessor:
    """Handles audio processing for the Audio Volume Normalizer."""
    
    def __init__(self):
        """Initialize the AudioProcessor."""
        self.loudness_cache: Dict[str, float] = {}
        self.encoding_cache: Dict[str, Dict[str, any]] = {}
        self.max_loudness: float = 0.0
        self.reference_file: Optional[str] = None
        self.reference_encoding: Optional[Dict[str, any]] = None
    
    def load_audio(self, file_path: str) -> Optional[AudioSegment]:
        """
        Load an audio file (MP3 or WAV) into a pydub AudioSegment.
        
        Args:
            file_path: Path to the audio file
            
        Returns:
            AudioSegment object or None if loading fails
            
        Raises:
            FileNotFoundError: If the file doesn't exist
        """
        if not os.path.exists(file_path):
            logger.error(f"File not found: {file_path}")
            raise FileNotFoundError(f"File not found: {file_path}")
        
        try:
            logger.debug(f"Loading audio file: {file_path}")
            
            # Determine file type and load accordingly
            if file_path.lower().endswith('.mp3'):
                audio = AudioSegment.from_mp3(file_path)
            elif file_path.lower().endswith('.wav'):
                audio = AudioSegment.from_wav(file_path)
            else:
                logger.error(f"Unsupported file format: {file_path}")
                return None
                
            return audio
        except Exception as e:
            # Log the error but don't crash the application
            logger.error(f"Failed to load audio file {file_path}: {e}")
            return None
    
    def calculate_loudness(self, audio: AudioSegment) -> float:
        """
        Calculate the loudness of an audio segment.
        
        This uses the RMS (Root Mean Square) method to calculate loudness,
        which is a common way to measure the average power of an audio signal.
        
        Args:
            audio: pydub AudioSegment
            
        Returns:
            Loudness value (higher is louder)
        """
        # Convert audio to numpy array for efficient processing
        samples = np.array(audio.get_array_of_samples())
        
        # Calculate RMS (Root Mean Square)
        rms = np.sqrt(np.mean(np.square(samples.astype(np.float64))))
        
        # Convert to dB (decibels)
        if rms > 0:
            db = 20 * np.log10(rms)
        else:
            db = -96.0  # Approximate silence level
        
        return db
    
    def get_encoding_parameters(self, audio: AudioSegment) -> Dict[str, any]:
        """
        Extract encoding parameters from an audio segment.
        
        Args:
            audio: pydub AudioSegment
            
        Returns:
            Dictionary containing encoding parameters
        """
        return {
            "sample_rate": audio.frame_rate,
            "channels": audio.channels,
            "bits_per_sample": audio.sample_width * 8,
            "format": "mp3" if hasattr(audio, "_data") and audio.channels > 0 else "wav"
        }
    
    def analyze_file(self, file_path: str) -> Optional[float]:
        """
        Analyze an audio file to determine its loudness and encoding parameters.
        
        Args:
            file_path: Path to the audio file
            
        Returns:
            Loudness value or None if analysis fails
        """
        try:
            audio = self.load_audio(file_path)
            if audio is None:
                return None
            
            # Get encoding parameters
            encoding = self.get_encoding_parameters(audio)
            self.encoding_cache[file_path] = encoding
            
            # Calculate loudness
            loudness = self.calculate_loudness(audio)
            self.loudness_cache[file_path] = loudness
            
            # Update max loudness if this file is louder
            if loudness > self.max_loudness:
                self.max_loudness = loudness
                self.reference_file = file_path
                self.reference_encoding = encoding
                logger.debug(f"New reference file: {file_path} (loudness: {loudness:.2f} dB)")
                logger.debug(f"Reference encoding: {self.reference_encoding}")
            
            return loudness
        except Exception as e:
            logger.error(f"Error analyzing file {file_path}: {e}")
            return None
    
    def analyze_files(
        self, 
        file_paths: List[str], 
        progress_callback: Optional[Callable[[int, int, str], None]] = None
    ) -> Tuple[Dict[str, float], Dict[str, Dict[str, any]]]:
        """
        Analyze multiple audio files to determine their loudness and encoding parameters.
        
        Args:
            file_paths: List of paths to audio files (MP3 and WAV)
            progress_callback: Optional callback function to report progress
                               Args: (current_index, total_count, current_file)
            
        Returns:
            Tuple of (loudness_cache, encoding_cache):
                - loudness_cache: Dictionary mapping file paths to loudness values
                - encoding_cache: Dictionary mapping file paths to encoding parameters
        """
        self.loudness_cache = {}
        self.encoding_cache = {}
        self.max_loudness = -96.0  # Start with silence level
        self.reference_file = None
        self.reference_encoding = None
        
        total_files = len(file_paths)
        
        for i, file_path in enumerate(file_paths):
            if progress_callback:
                progress_callback(i, total_files, file_path)
            
            loudness = self.analyze_file(file_path)
            if loudness is not None:
                logger.debug(f"File {file_path}: loudness = {loudness:.2f} dB")
        
        if progress_callback:
            progress_callback(total_files, total_files, "Analysis complete")
        
        if self.reference_file:
            logger.info(f"Reference file: {self.reference_file} (loudness: {self.max_loudness:.2f} dB)")
        else:
            logger.warning("No valid reference file found")
        
        return (self.loudness_cache, self.encoding_cache)
    
    def normalize_file(
        self, 
        file_path: str,
        target_loudness: Optional[float] = None
    ) -> Optional[AudioSegment]:
        """
        Normalize the volume of an audio file.
        
        Args:
            file_path: Path to the audio file (MP3 or WAV)
            target_loudness: Target loudness in dB (if None, uses max_loudness)
            
        Returns:
            Normalized AudioSegment or None if normalization fails
        """
        if target_loudness is None:
            if self.max_loudness == 0:
                logger.error("No reference loudness available")
                return None
            target_loudness = self.max_loudness
        
        try:
            # Load the audio file
            audio = self.load_audio(file_path)
            if audio is None:
                return None
            
            # Get current loudness
            current_loudness = self.loudness_cache.get(file_path)
            if current_loudness is None:
                current_loudness = self.calculate_loudness(audio)
            
            # Calculate the gain needed
            gain_db = target_loudness - current_loudness
            
            # Apply the gain
            if abs(gain_db) > 0.1:  # Only apply if the difference is significant
                logger.debug(f"Normalizing {file_path}: applying {gain_db:.2f} dB gain")
                normalized_audio = audio.apply_gain(gain_db)
                return normalized_audio
            else:
                logger.debug(f"File {file_path} already at target loudness")
                return audio
        except Exception as e:
            logger.error(f"Error normalizing file {file_path}: {e}")
            return None
    
    def normalize_files(
        self,
        file_paths: List[str],
        file_manager: any,  # Avoid circular import
        progress_callback: Optional[Callable[[int, int, str], None]] = None
    ) -> Tuple[int, int]:
        """
        Normalize multiple audio files to match the loudest file.
        
        Args:
            file_paths: List of paths to audio files (MP3 and WAV)
            file_manager: FileManager instance for saving files
            progress_callback: Optional callback function to report progress
                               Args: (current_index, total_count, current_file)
            
        Returns:
            Tuple of (number of successful normalizations, number of failed normalizations)
        """
        if not self.reference_file:
            logger.error("No reference file available for normalization")
            return (0, 0)
        
        total_files = len(file_paths)
        success_count = 0
        fail_count = 0
        
        for i, file_path in enumerate(file_paths):
            if progress_callback:
                progress_callback(i, total_files, file_path)
            
            # Skip the reference file if it's already at max loudness
            if file_path == self.reference_file:
                logger.debug(f"Skipping reference file: {file_path}")
                success_count += 1
                continue
            
            # Normalize the file
            normalized_audio = self.normalize_file(file_path)
            if normalized_audio is None:
                logger.error(f"Failed to normalize {file_path}")
                fail_count += 1
                continue
            
            # Save the normalized file with reference encoding
            if file_manager.save_normalized_file(file_path, normalized_audio, self.reference_encoding):
                success_count += 1
            else:
                fail_count += 1
        
        if progress_callback:
            progress_callback(total_files, total_files, "Normalization complete")
        
        logger.info(f"Normalized {success_count} files, {fail_count} failed")
        return (success_count, fail_count)