import numpy as np
from scipy.signal import butter, sosfilt
from .base import BaseProcessorModule

class MsLimiterModule(BaseProcessorModule):
    """Surgical Mid-Side Limiter - Updated for High-End Transparency."""
    
    def process(self, audio: np.ndarray, side_limit_db: float = -3.0, mid_gain_db: float = 0.0) -> np.ndarray:
        """Limit the side-channel energy only in the problem range (below 2kHz)."""
        if audio.ndim == 1: return audio
        
        sr = self.sample_rate
        # Convert to MS
        mid = 0.5 * (audio[:, 0] + audio[:, 1])
        side = 0.5 * (audio[:, 0] - audio[:, 1])
        
        # 1. Split Side channel into Problem (Low) and Transparent (High)
        # We only want to stabilize up to 2kHz
        sos_lp = butter(4, 2000, 'lp', fs=sr, output='sos')
        sos_hp = butter(4, 2000, 'hp', fs=sr, output='sos')
        
        side_low = sosfilt(sos_lp, side)
        side_high = sosfilt(sos_hp, side) # THE 'AIR' (Protect this!)
        
        # 2. Limit the LOW SIDES only
        win_size = int(sr * 0.05)
        envelope = np.convolve(np.abs(side_low), np.ones(win_size)/win_size, mode='same')
        thresh_lin = 10 ** (side_limit_db / 20)
        
        reduction = np.ones_like(envelope)
        mask = envelope > thresh_lin
        reduction[mask] = thresh_lin / (envelope[mask] + 1e-9)
        
        side_low_limited = side_low * reduction
        
        # 3. Reconstruct Side (Limited Lows + Untouched Highs)
        side_reconstructed = side_low_limited + side_high
        
        # 4. Apply Mid Gain
        mid_boosted = mid * (10 ** (mid_gain_db / 20))
        
        # Convert back to LR
        left = mid_boosted + side_reconstructed
        right = mid_boosted - side_reconstructed
        
        return np.column_stack((left, right))