"""Unit tests for the IR convolution engine (IRLoader).

Each test validates a specific aspect of FFT overlap-add convolution:
synthetic IR identity, silence handling, wet/dry mix, toggle, state
management across blocks, directory listing, and performance budget.

All tests use 256-sample blocks at 48kHz to match real-time operation.
"""

from __future__ import annotations

import itertools
import time
from pathlib import Path

import numpy as np
import pytest

from src.dsp.ir_loader import IRLoader, IRFile, _next_pow2

# ── Test constants ──────────────────────────────────────────────────

BLOCK_SIZE = 256
SAMPLE_RATE = 48000

SILENCE = np.zeros(BLOCK_SIZE, dtype=np.float32)
SINE_TONE = (np.sin(2 * np.pi * 440.0 * np.arange(BLOCK_SIZE) / SAMPLE_RATE)
             .astype(np.float32))
HALF_SCALE = np.full(BLOCK_SIZE, 0.5, dtype=np.float32)
FULL_SCALE = np.full(BLOCK_SIZE, 0.99, dtype=np.float32)

# A synthetic IR that acts as a band-pass filter (simple chirp)
_SHORT_IR = (
    np.sin(2 * np.pi * 500.0 * np.arange(256) / SAMPLE_RATE)
    * np.exp(-np.arange(256) / 64.0)
).astype(np.float32)

_MEDIUM_IR = (
    np.sin(2 * np.pi * 800.0 * np.arange(1024) / SAMPLE_RATE)
    * np.exp(-np.arange(1024) / 128.0)
).astype(np.float32)

_LONG_IR = (
    np.sin(2 * np.pi * 400.0 * np.arange(4096) / SAMPLE_RATE)
    * np.exp(-np.arange(4096) / 512.0)
).astype(np.float32)


# ── Helpers ─────────────────────────────────────────────────────────

def _load_synthetic(ir: IRLoader, ir_data: np.ndarray) -> bool:
    """Load a synthetic IR by writing a temp .wav file and loading it."""
    from scipy.io import wavfile
    import tempfile
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    wavfile.write(tmp.name, SAMPLE_RATE, ir_data)
    result = ir.load_ir(tmp.name)
    Path(tmp.name).unlink()
    return result


# ═══════════════════════════════════════════════════════════════════
# 1. Basic IR loading
# ═══════════════════════════════════════════════════════════════════

class TestIRLoading:
    def test_load_short_ir(self):
        """Load a 256-tap synthetic IR successfully."""
        ir = IRLoader()
        assert _load_synthetic(ir, _SHORT_IR), "Should load successfully"
        assert ir.is_loaded, "is_loaded should be True"
        assert ir.current_ir is not None
        assert ir.current_ir.num_taps == 256

    def test_load_medium_ir(self):
        """Load a 1024-tap IR."""
        ir = IRLoader()
        assert _load_synthetic(ir, _MEDIUM_IR)
        assert ir.current_ir is not None
        assert ir.current_ir.num_taps == 1024

    def test_load_long_ir(self):
        """Load a 4096-tap IR (long cabinet)."""
        ir = IRLoader()
        assert _load_synthetic(ir, _LONG_IR)
        assert ir.current_ir is not None
        assert ir.current_ir.num_taps == 4096

    def test_load_max_taps(self):
        """8192-tap IR should load (the max)."""
        long = (
            np.sin(2 * np.pi * 200.0 * np.arange(8192) / SAMPLE_RATE)
            * np.exp(-np.arange(8192) / 1024.0)
        ).astype(np.float32)
        ir = IRLoader()
        assert _load_synthetic(ir, long)
        assert ir.current_ir is not None
        assert ir.current_ir.num_taps == 8192

    def test_load_nonexistent_file(self):
        """Non-existent file returns False."""
        ir = IRLoader()
        assert not ir.load_ir("/nonexistent/cab.wav"), "Should fail"

    def test_load_non_wav(self):
        """Non-.wav file returns False."""
        import tempfile
        tmp = tempfile.NamedTemporaryFile(suffix=".txt", delete=False)
        tmp.write(b"hello")
        tmp.close()
        ir = IRLoader()
        assert not ir.load_ir(tmp.name), "Should reject non-wav"
        Path(tmp.name).unlink()

    def test_unload(self):
        """Unload clears all state."""
        ir = IRLoader()
        _load_synthetic(ir, _SHORT_IR)
        ir.unload()
        assert not ir.is_loaded
        assert ir.current_ir is None

    def test_metadata_correct(self):
        """IRFile metadata reflects actual file content."""
        ir = IRLoader()
        _load_synthetic(ir, _MEDIUM_IR)
        assert ir.current_ir is not None
        assert ir.current_ir.sample_rate == SAMPLE_RATE
        expected_ms = (1024 / SAMPLE_RATE) * 1000
        assert abs(ir.current_ir.length_ms - expected_ms) < 0.1


# ═══════════════════════════════════════════════════════════════════
# 2. FFT overlap-add convolution (correctness)
# ═══════════════════════════════════════════════════════════════════

class TestConvolution:
    def test_silence_in_silence_out(self):
        """Silence input produces silence output."""
        ir = IRLoader()
        _load_synthetic(ir, _SHORT_IR)
        out = ir.process(SILENCE)
        assert np.max(np.abs(out)) == 0.0, "Silence in → silence out"

    def test_identity_with_dirac(self):
        """Convolving with a Dirac impulse (first sample = 1) returns input."""
        dirac = np.zeros(256, dtype=np.float32)
        dirac[0] = 1.0
        ir = IRLoader()
        _load_synthetic(ir, dirac)
        out = ir.process(SINE_TONE * 0.5)
        # Allow small error due to FFT floating point
        assert np.allclose(out, SINE_TONE * 0.5, atol=1e-5), \
            "Dirac IR should act as identity"

    def test_amplitude_scaling(self):
        """Convolving with a scaled Dirac scales output by the same factor."""
        dirac = np.zeros(256, dtype=np.float32)
        dirac[0] = 0.5  # half amplitude
        ir = IRLoader()
        _load_synthetic(ir, dirac)
        out = ir.process(SINE_TONE * 0.5)
        assert np.allclose(out, SINE_TONE * 0.25, atol=1e-5), \
            "0.5 Dirac should scale amplitude 0.5x"

    def test_output_length_matches_input(self):
        """process() returns a block of the same length as input."""
        ir = IRLoader()
        _load_synthetic(ir, _MEDIUM_IR)
        out = ir.process(SINE_TONE)
        assert len(out) == len(SINE_TONE), \
            f"Output length {len(out)} should equal input {len(SINE_TONE)}"

    def test_output_range(self):
        """Processed output stays in [-1, 1]."""
        ir = IRLoader()
        _load_synthetic(ir, _LONG_IR)
        out = ir.process(FULL_SCALE)
        assert np.all(out >= -1.0) and np.all(out <= 1.0), \
            "Output must be clipped to [-1, 1]"

    def test_no_nan_or_inf(self):
        """No NaN/Inf in output for any reasonable input."""
        ir = IRLoader()
        _load_synthetic(ir, _MEDIUM_IR)
        out = ir.process(SINE_TONE * 0.7)
        assert np.all(np.isfinite(out)), "Output must be finite"

    def test_lazy_fft_recompute_on_first_block(self):
        """FFT is computed on first process() call, not at load time."""
        ir = IRLoader()
        _load_synthetic(ir, _SHORT_IR)
        # FFT not computed yet
        assert ir._conv_fft_len == 0, "FFT should not be pre-computed"
        out = ir.process(HALF_SCALE)
        assert ir._conv_fft_len > 0, "FFT should be computed after first process"
        assert len(out) == BLOCK_SIZE


# ═══════════════════════════════════════════════════════════════════
# 3. Overlap-add state across blocks
# ═══════════════════════════════════════════════════════════════════

class TestOverlapAdd:
    def test_tail_propagates(self):
        """Convolution tail from block N carries into block N+1.

        With a long IR, a single impulse should produce output that
        spans multiple blocks.
        """
        ir_len = 512
        long_ir = (
            np.sin(2 * np.pi * 600.0 * np.arange(ir_len) / SAMPLE_RATE)
            * np.exp(-np.arange(ir_len) / 64.0)
        ).astype(np.float32)

        ir = IRLoader()
        _load_synthetic(ir, long_ir)

        # Send one block with a single impulse
        impulse = np.zeros(BLOCK_SIZE, dtype=np.float32)
        impulse[0] = 1.0

        out1 = ir.process(impulse)
        # First block should have energy from convolution
        assert np.max(np.abs(out1)) > 0.1, "First block should have output"

        # Second block with silence should still have tail energy
        out2 = ir.process(SILENCE)
        if np.max(np.abs(out2)) == 0.0:
            # IR shorter than block — no tail. Acceptable.
            pass
        else:
            # There is a tail — should decay
            out3 = ir.process(SILENCE)
            assert np.max(np.abs(out3)) <= np.max(np.abs(out2)) + 0.001, \
                "Tail should not increase"

    def test_consecutive_blocks_differ(self):
        """Consecutive identical input blocks produce different output
        when IR is longer than block (overlap-add state changes)."""
        ir = IRLoader()
        # IR longer than block ensures overlap state
        ir_data = (
            np.sin(2 * np.pi * 300.0 * np.arange(1024) / SAMPLE_RATE)
            * np.exp(-np.arange(1024) / 256.0)
        ).astype(np.float32)
        _load_synthetic(ir, ir_data)

        out1 = ir.process(SINE_TONE)
        out2 = ir.process(SINE_TONE)
        # If IR length > block, the first and second blocks should differ
        # because the second block convolves with the existing tail
        assert not np.allclose(out1, out2, atol=1e-4), \
            "Consecutive blocks should differ with overlap-add"

    def test_reset_tail(self):
        """reset_tail() clears the overlap state."""
        ir = IRLoader()
        ir_data = (
            np.sin(2 * np.pi * 300.0 * np.arange(1024) / SAMPLE_RATE)
            * np.exp(-np.arange(1024) / 256.0)
        ).astype(np.float32)
        _load_synthetic(ir, ir_data)

        # Fill overlap buffer
        ir.process(FULL_SCALE)
        ir.process(FULL_SCALE)
        ir.process(FULL_SCALE)

        tail_before = ir._tail.copy()
        ir.reset_tail()
        assert len(ir._tail) == 0, "Tail should be empty after reset"
        # And subsequent process with silence should be silent
        out = ir.process(SILENCE)
        assert np.max(np.abs(out)) == 0.0, "Silence after tail reset"

    def test_disable_clears_tail(self):
        """Disabling the IR clears the tail buffer."""
        ir = IRLoader()
        _load_synthetic(ir, _MEDIUM_IR)
        ir.process(FULL_SCALE)
        ir.process(FULL_SCALE)
        ir.enabled = False
        out = ir.process(SILENCE)
        assert np.max(np.abs(out)) == 0.0, "Disabled IR should pass silence"
        ir.enabled = True
        # Should start clean
        out = ir.process(SILENCE)
        assert np.max(np.abs(out)) == 0.0, "Re-enabled IR with silence"


# ═══════════════════════════════════════════════════════════════════
# 4. Wet/dry mix control
# ═══════════════════════════════════════════════════════════════════

class TestMix:
    def test_dry_only_bypass(self):
        """100% dry = original signal unchanged (no convolution)."""
        ir = IRLoader()
        _load_synthetic(ir, _MEDIUM_IR)
        ir.set_mix(wet=0.0, dry=1.0)
        out = ir.process(SINE_TONE)
        assert np.allclose(out, SINE_TONE, atol=1e-5), \
            "100% dry should pass through original"

    def test_wet_only_full_convolution(self):
        """100% wet = fully convolved signal."""
        ir = IRLoader()
        _load_synthetic(ir, _SHORT_IR)
        ir.set_mix(wet=1.0, dry=0.0)
        out_wet = ir.process(SINE_TONE)
        ir2 = IRLoader()
        _load_synthetic(ir2, _SHORT_IR)
        out_default = ir2.process(SINE_TONE)
        assert np.allclose(out_wet, out_default, atol=1e-5), \
            "Default mix (1.0/0.0) should equal explicit 100% wet"

    def test_balanced_mix(self):
        """50/50 mix produces mid-way output."""
        ir = IRLoader()
        _load_synthetic(ir, _SHORT_IR)
        ir.set_mix(wet=0.5, dry=0.5)
        out = ir.process(HALF_SCALE)
        assert np.max(np.abs(out)) > 0, "50/50 mix should produce output"
        assert np.all(out >= -1.0) and np.all(out <= 1.0)

    def test_wet_property_setter(self):
        """wet property setter works."""
        ir = IRLoader()
        assert ir.wet == 1.0, "Default wet should be 1.0"
        ir.wet = 0.3
        assert ir.wet == 0.3
        ir.wet = 1.5  # Clamp
        assert ir.wet == 1.0

    def test_dry_property_setter(self):
        """dry property setter works."""
        ir = IRLoader()
        assert ir.dry == 0.0, "Default dry should be 0.0"
        ir.dry = 0.7
        assert ir.dry == 0.7
        ir.dry = -0.5  # Clamp
        assert ir.dry == 0.0


# ═══════════════════════════════════════════════════════════════════
# 5. Enable/disable toggle
# ═══════════════════════════════════════════════════════════════════

class TestToggle:
    def test_disabled_passes_dry(self):
        """When disabled, process() returns input unchanged."""
        ir = IRLoader()
        _load_synthetic(ir, _MEDIUM_IR)
        ir.enabled = False
        out = ir.process(SINE_TONE)
        assert np.allclose(out, SINE_TONE, atol=1e-5), \
            "Disabled IR should pass-through"

    def test_disabled_does_not_convolution(self):
        """Disabled IR should have no convolution artifacts."""
        ir = IRLoader()
        _load_synthetic(ir, _SHORT_IR)
        ir.enabled = False
        out = ir.process(FULL_SCALE)
        assert np.allclose(out, FULL_SCALE, atol=1e-5), \
            "Disabled: full-scale should pass through unchanged"

    def test_toggle_recovers(self):
        """Toggle off then on recovers convolution."""
        ir = IRLoader()
        _load_synthetic(ir, _SHORT_IR)
        ir.enabled = False
        ir.process(FULL_SCALE)  # Should be pass-through
        ir.enabled = True
        out = ir.process(FULL_SCALE)
        assert not np.allclose(out, FULL_SCALE, atol=1e-2), \
            "Re-enabled IR should convolve (shape differs)"

    def test_default_enabled(self):
        """IRLoader starts enabled."""
        ir = IRLoader()
        assert ir.enabled, "Default state should be enabled"


# ═══════════════════════════════════════════════════════════════════
# 6. Directory listing
# ═══════════════════════════════════════════════════════════════════

class TestDirectoryListing:
    def test_empty_dir(self, tmp_path):
        """Empty IR directory returns empty list."""
        ir = IRLoader(tmp_path)
        irs = ir.get_irs()
        assert len(irs) == 0, "Empty dir should return []"

    def test_finds_wav_files(self, tmp_path):
        """get_irs() finds .wav files in the IR directory."""
        from scipy.io import wavfile
        # Write a .wav
        wavfile.write(str(tmp_path / "test_ir.wav"), SAMPLE_RATE, _SHORT_IR)
        ir = IRLoader(tmp_path)
        irs = ir.get_irs()
        assert len(irs) == 1
        assert irs[0].name == "test_ir"
        assert irs[0].num_taps == 256

    def test_skips_non_wav(self, tmp_path):
        """Non-.wav files are skipped."""
        from scipy.io import wavfile
        wavfile.write(str(tmp_path / "good.wav"), SAMPLE_RATE, _SHORT_IR)
        (tmp_path / "not_an_ir.txt").write_text("hello")
        ir = IRLoader(tmp_path)
        irs = ir.get_irs()
        assert len(irs) == 1
        assert irs[0].name == "good"

    def test_returns_sorted(self, tmp_path):
        """get_irs() returns files in sorted order."""
        from scipy.io import wavfile
        wavfile.write(str(tmp_path / "b.wav"), SAMPLE_RATE, _SHORT_IR)
        wavfile.write(str(tmp_path / "a.wav"), SAMPLE_RATE, _SHORT_IR)
        ir = IRLoader(tmp_path)
        irs = ir.get_irs()
        assert [ir.name for ir in irs] == ["a", "b"]


# ═══════════════════════════════════════════════════════════════════
# 7. Performance budget < 5ms per block
# ═══════════════════════════════════════════════════════════════════

class TestPerformance:
    def test_short_ir_under_budget(self):
        """256-tap IR processes in < 5ms."""
        ir = IRLoader()
        _load_synthetic(ir, _SHORT_IR)
        # Warm up — first block computes FFT
        ir.process(HALF_SCALE)
        # Time a few blocks
        times = []
        for _ in range(10):
            start = time.perf_counter()
            ir.process(HALF_SCALE)
            elapsed = (time.perf_counter() - start) * 1000  # ms
            times.append(elapsed)
        mean_ms = sum(times) / len(times)
        assert mean_ms < 5.0, \
            f"Short IR: {mean_ms:.2f}ms avg, expected < 5ms"

    def test_medium_ir_under_budget(self):
        """1024-tap IR processes in < 5ms."""
        ir = IRLoader()
        _load_synthetic(ir, _MEDIUM_IR)
        ir.process(HALF_SCALE)  # warm
        times = []
        for _ in range(10):
            start = time.perf_counter()
            ir.process(HALF_SCALE)
            elapsed = (time.perf_counter() - start) * 1000
            times.append(elapsed)
        mean_ms = sum(times) / len(times)
        assert mean_ms < 5.0, \
            f"Medium IR: {mean_ms:.2f}ms avg, expected < 5ms"

    def test_long_ir_under_budget(self):
        """4096-tap IR processes in < 5ms."""
        ir = IRLoader()
        _load_synthetic(ir, _LONG_IR)
        ir.process(HALF_SCALE)  # warm
        times = []
        for _ in range(10):
            start = time.perf_counter()
            ir.process(HALF_SCALE)
            elapsed = (time.perf_counter() - start) * 1000
            times.append(elapsed)
        mean_ms = sum(times) / len(times)
        assert mean_ms < 5.0, \
            f"Long IR: {mean_ms:.2f}ms avg, expected < 5ms"

    def test_max_taps_under_budget(self):
        """8192-tap IR processes in < 5ms."""
        max_ir = (
            np.sin(2 * np.pi * 200.0 * np.arange(8192) / SAMPLE_RATE)
            * np.exp(-np.arange(8192) / 1024.0)
        ).astype(np.float32)
        ir = IRLoader()
        _load_synthetic(ir, max_ir)
        ir.process(HALF_SCALE)  # warm
        times = []
        for _ in range(10):
            start = time.perf_counter()
            ir.process(HALF_SCALE)
            elapsed = (time.perf_counter() - start) * 1000
            times.append(elapsed)
        mean_ms = sum(times) / len(times)
        assert mean_ms < 5.0, \
            f"Max IR: {mean_ms:.2f}ms avg, expected < 5ms"


# ═══════════════════════════════════════════════════════════════════
# 8. Edge cases
# ═══════════════════════════════════════════════════════════════════

class TestEdgeCases:
    def test_process_before_load(self):
        """process() with no IR loaded returns input unchanged."""
        ir = IRLoader()
        out = ir.process(SINE_TONE)
        assert np.allclose(out, SINE_TONE), \
            "No IR loaded = passthrough"

    def test_process_with_tiny_ir(self):
        """IR shorter than block size works correctly."""
        tiny_ir = np.array([0.5, 0.3], dtype=np.float32)
        ir = IRLoader()
        _load_synthetic(ir, tiny_ir)
        out = ir.process(SINE_TONE)
        assert np.all(np.isfinite(out))
        assert np.all(out >= -1.0) and np.all(out <= 1.0)

    def test_load_ir_after_unload(self):
        """Load-then-unload-then-reload cycle works."""
        ir = IRLoader()
        _load_synthetic(ir, _SHORT_IR)
        ir.unload()
        _load_synthetic(ir, _MEDIUM_IR)
        assert ir.is_loaded
        assert ir.current_ir is not None
        assert ir.current_ir.num_taps == 1024

    def test_many_consecutive_blocks_no_drift(self):
        """100 consecutive blocks should not clip/drift/clog."""
        ir = IRLoader()
        _load_synthetic(ir, _MEDIUM_IR)
        for i in range(100):
            out = ir.process(SINE_TONE)
            assert np.all(np.isfinite(out)), f"NaN at block {i}"
            assert np.all(out >= -1.0) and np.all(out <= 1.0), \
                f"Clip violation at block {i}"

    def test_single_sample_block(self):
        """Process a single-sample block without error."""
        ir = IRLoader()
        _load_synthetic(ir, _SHORT_IR)
        block = np.array([0.5], dtype=np.float32)
        out = ir.process(block)
        assert len(out) == 1
        assert np.all(np.isfinite(out))

    def test_int16_normalisation(self):
        """WAV int16 data normalises to float32 [-1, 1]."""
        from scipy.io import wavfile
        import tempfile
        int16_data = (np.arange(256) - 128).astype(np.int16) * 256
        tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
        wavfile.write(tmp.name, SAMPLE_RATE, int16_data)
        ir = IRLoader()
        ir.load_ir(tmp.name)
        Path(tmp.name).unlink()
        assert ir._ir_data is not None
        assert ir._ir_data.dtype == np.float32
        assert np.max(np.abs(ir._ir_data)) <= 1.0 + 1e-5, \
            "Normalised float32 should be in [-1, 1]"


# ═══════════════════════════════════════════════════════════════════
# 9. _next_pow2 utility
# ═══════════════════════════════════════════════════════════════════

class TestNextPow2:
    def test_exact_pow2(self):
        assert _next_pow2(1024) == 1024
        assert _next_pow2(1) == 1
        assert _next_pow2(2) == 2

    def test_rounds_up(self):
        assert _next_pow2(3) == 4
        assert _next_pow2(5) == 8
        assert _next_pow2(100) == 128

    def test_large_number(self):
        assert _next_pow2(8447) == 16384  # typical IR FFT size
        assert _next_pow2(16383) == 16384

    def test_zero(self):
        # _next_pow2(0) = 1 (1 << -1? No: (0-1).bit_length() = 0, 1<<0 = 1)
        # For our use case, n is always >= 1, but just in case:
        assert _next_pow2(1) == 1