import numpy as np
from scipy.signal import lfilter
from .base import BaseProcessorModule

class TransientModule(BaseProcessorModule):
    """Transient Shaper Module."""
    
    def process(self, audio: np.ndarray, boost_db: float = 3.0, sustain_db: float = -2.0) -> np.ndarray:
        sr = self.sample_rate
        
        def process_channel(channel):
            # Fast envelope (Attack)
            att_samples = int(sr * 0.005) # 5ms
            alpha_att = np.exp(-1.0 / att_samples)
            env_fast = lfilter([1.0 - alpha_att], [1.0, -alpha_att], np.abs(channel))
            
            # Slow envelope (Sustain)
            rel_samples = int(sr * 0.050) # 50ms
            alpha_rel = np.exp(-1.0 / rel_samples)
            env_slow = lfilter([1.0 - alpha_rel], [1.0, -alpha_rel], np.abs(channel))
            
            # Transient detection (Difference)
            transient = env_fast - env_slow
            
            # Gains
            gain_att = 10 ** (boost_db / 20)
            gain_sus = 10 ** (sustain_db / 20)
            
            # Apply
            # If transient > 0, it's an attack phase -> boost
            # If transient < 0, it's a sustain phase -> cut/boost
            
            # Ratio mask
            ratio = np.clip(transient, -1.0, 1.0) 
            
            # This is a simplified shaper. 
            # Real ones use the ratio to modulate gain.
            # Here we'll just add the difference back scaled.
            
            # Simpler approach:
            # Output = Original + (Transient * Boost) + (Sustain * Gain)
            # But that's not quite right.
            
            # Standard approach:
            # Gain = 1.0 + (Transient * BoostFactor)
            gain_curve = np.ones_like(channel)
            mask_att = transient > 0
            gain_curve[mask_att] = 1.0 + (transient[mask_att] * (gain_att - 1.0))
            gain_curve[~mask_att] = 1.0 + (transient[~mask_att] * (1.0 - gain_sus)) # Invert logic for sustain
            
            return channel * gain_curve

        if audio.ndim == 1: return process_channel(audio)
        else:
            left = process_channel(audio[:, 0])
            right = process_channel(audio[:, 1])
            return np.column_stack((left, right))