import argparse
import warnings
import numpy as np
from pathlib import Path
from scipy.io import wavfile
from scipy.io.wavfile import WavFileWarning
from scipy.signal import fftconvolve
from scipy.signal.windows import tukey

# Hide metadata warnings emitted by scipy when reading non-standard WAV headers
warnings.filterwarnings('ignore', category=WavFileWarning)


def load_wav_normalized(filepath):
    """
    Loads a WAV file and strictly normalizes its audio data to float64 
    in the range [-1.0, +1.0] based on the original integer or float format.
    """
    sr, raw_data = wavfile.read(filepath)
    original_dtype = raw_data.dtype
    data = raw_data.astype(np.float64)

    # Scale integer PCM formats based on original dtype BEFORE casting side-effects
    if np.issubdtype(original_dtype, np.int16):
        data /= 32768.0
    elif np.issubdtype(original_dtype, np.int32):
        data /= 2147483648.0
    elif np.issubdtype(original_dtype, np.uint8):
        data = (data - 128.0) / 128.0
    # Floating-point WAVs (float32/float64) are already in [-1.0, 1.0] range

    return sr, data


def apply_tukey_window(data: np.ndarray, alpha: float) -> np.ndarray:
    """
    Applies a Tukey (tapered cosine) window to an impulse/audio signal.
    
    Parameters:
        data (np.ndarray): Input 1D (mono) or 2D (multichannel) audio array.
        alpha (float): Shape parameter (0.0 <= alpha <= 1.0).
                       e.g., 0.05 = 5% total tapered cosine fade (2.5% each edge).
    Returns:
        np.ndarray: Windowed audio data.
    """
    if alpha <= 0.0:
        return data  # Rectangular window (no-op)

    num_taps = len(data)
    window = tukey(num_taps, alpha=alpha, sym=True)

    # Adjust window shape for 2D multichannel arrays if needed
    if data.ndim == 2:
        window = window[:, np.newaxis]

    return data * window


def main():
    parser = argparse.ArgumentParser(description="Cascade FIR WAV files clean and normalized")
    # Optional --gain argument with default value 0.0 dB
    parser.add_argument('--gain', '-g', type=float, default=0.0,
        help="Gain adjustment in dB (default: 0.0 dB)"
    )
    
    # Optional --alpha argument with range validation and default value 0.0
    parser.add_argument('--alpha', '-a', type=float, default=0.0, 
        help="Tukey shape parameter between 0.0 and 1.0 (default: 0.0 = 0%% total fade)"
    )

    parser.add_argument('Output_file', help="Path to output WAV file (e.g., output.wav)")

    # Positional argument for one or more input files
    parser.add_argument('Input_files', nargs='+', help="input1.wav input2.wav ...")

    args = parser.parse_args()

    # Validate alpha parameter bounds
    if not (0.0 <= args.alpha <= 1.0):
        parser.error(f"Tukey alpha must be between 0.0 and 1.0 (got {args.alpha})")

    # Validate input files
    missing_files = [f for f in args.Input_files if not Path(f).is_file()]
    if missing_files:
        parser.error(f"File(s) not found: {', '.join(missing_files)}")

    print(f"\nFIR-cascade: {len(args.Input_files)} file(s) | Gain: {args.gain:+.2f} dB | Tukey alpha: {args.alpha}")

    # 1. Load the first FIR impulse response
    sr_base, data = load_wav_normalized(args.Input_files[0])
    print(f" -> Loaded:          {args.Input_files[0]} ({len(data)} Taps, {sr_base} Hz)")

    # 2. Sequentially convolve with subsequent FIR filter files
    if len(args.Input_files) > 1:
        for f in args.Input_files[1:]:
            sr, h = load_wav_normalized(f)
            if sr != sr_base:
                raise ValueError(f"Sample rate mismatch in {f} ({sr} Hz vs {sr_base} Hz)!")

            print(f" -> Convolving with: {f} ({len(h)} Taps)")

            # Handle mono vs stereo/multichannel broadcasting alignment
            if data.ndim == 1 and h.ndim == 2:
                data = data[:, np.newaxis]
            elif data.ndim == 2 and h.ndim == 1:
                h = h[:, np.newaxis]

            # Perform full linear convolution along the time axis (axis 0)
            data = fftconvolve(data, h, mode='full', axes=0)

    # 3. Apply user-defined gain adjustment
    linear_gain = 10.0 ** (args.gain / 20.0)
    data = data * linear_gain

    # 4. Apply a Tukey (tapered cosine) window to ensure smooth edges 
    data = apply_tukey_window(data, alpha=args.alpha)

    # 5. Measure final peak level to warn against potential digital clipping
    peak_val = np.max(np.abs(data))
    peak_db = 20 * np.log10(peak_val) if peak_val > 0 else -140.0

    print(f"\nFinal Impulse Info:")
    print(f"    Output Length:   {len(data)} Taps ({len(data) / sr_base * 1000:.1f} ms)")
    print(f"    Peak level:      {peak_val:.6f} ({peak_db:.2f} dBFS)")

    if peak_val > 1.0:
        print("WARNING: Output peak exceeds 0 dBFS! Digital clipping will occur if played at full volume.")
        print(f"    --> Recommended Preamp reduction: -{peak_db:.2f} dB")

    # 6. Save cascaded impulse as 32-bit Float WAV to prevent internal rounding errors
    wavfile.write(args.Output_file, sr_base, data.astype(np.float32))
    print(f"Output saved to:     {args.Output_file}\n")


if __name__ == "__main__":
    main()