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

try:
    from ..pipeline.saturation import oversampled
except ImportError:
    from automaster_app.pipeline.saturation import oversampled


class VintageModule(BaseProcessorModule):
    """Surgical Vintage Saturation with Anti-Aliasing Guard."""
    
    def process(self, audio: np.ndarray, saturation: float = 0.0, character: str = "Tube", blend: float = 1.0) -> np.ndarray:
        if saturation <= 0: return audio
        
        sr = self.sample_rate
        original = audio.copy()
        drive = 1.0 + (saturation * 0.3)
        processed = audio * drive
        
        # 1. Anti-Aliasing Pre-Filter (Smooth the curve)
        if character == "Tube":
            processed = oversampled(
                lambda u: (np.arctan(u * 1.5) / 1.5) * 0.8 + (u * 0.2),
                processed,
            )
        elif character == "Solid":
            processed = oversampled(lambda u: np.tanh(u), processed)
        else:
            processed = oversampled(lambda u: np.tanh(u * 1.1) * 0.95, processed)
            
        # 2. Anti-Aliasing Post-Filter (Remove ultrasonic fold-back)
        # We cut at 18kHz inside the wet chain only
        sos = butter(4, 18000, 'lp', fs=sr, output='sos')
        if processed.ndim == 1:
            processed = sosfilt(sos, processed)
        else:
            processed[:, 0] = sosfilt(sos, processed[:, 0])
            processed[:, 1] = sosfilt(sos, processed[:, 1])
            
        # 3. Parallel Blend
        out = (processed * blend) + (original * (1.0 - blend))
        return np.clip(out, -1.0, 1.0)
