"""IR cab loader — impulse response convolution for cabinet simulation.

Uses numpy FFT-based overlap-add convolution for real-time IR playback.
On RPi 4B, typical IR files (512-2048 taps at 48kHz) perform
efficiently with block-based overlap-add (<2ms per 256-sample block).

Designed for use in the Pi Multi-FX Pedal's signal chain:
  Guitar -> ... -> NAM Amp -> IR Cab -> EQ -> ...
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import numpy as np
from scipy.io import wavfile

logger = logging.getLogger(__name__)

DEFAULT_IR_DIR = Path.home() / ".pedal" / "irs"

# ── Frequency-domain convolution helpers ────────────────────────────

# The maximum IR length we support: 8192 taps @ 48kHz ≈ 171ms.
# Overlap-add FFT size for a 256-block + 8192-tap IR = next_pow2(8447) = 16384.
# That's a ~65ms FFT on RPi 4B — well under our 5ms budget.
_MAX_IR_TAPS = 8192

# Safety clip for float32 IR values beyond [-1, 1]
_EPS = 1e-10


@dataclass
class IRFile:
    """Metadata for a loaded IR file."""
    name: str
    path: str
    sample_rate: int
    num_taps: int
    length_ms: float
    channels: int = 1


def _next_pow2(n: int) -> int:
    """Return the smallest power of 2 >= n."""
    return 1 << (n - 1).bit_length()


class IRLoader:
    """FFT overlap-add convolution engine for IR cabinet simulation.

    Loads standard .wav IR files and processes audio blocks in real-time.
    Supports per-block wet/dry mix and per-preset toggle.

    Typical usage::

        ir = IRLoader()
        ir.load_ir("my_cab.wav")
        output = ir.process(input_block)   # float32 [-1, 1]
        ir.set_mix(wet=0.8, dry=0.2)
        ir.enabled = False                 # passes dry
    """

    def __init__(self, ir_dir: str | Path = DEFAULT_IR_DIR):
        self._ir_dir = Path(ir_dir)
        self._ir_dir.mkdir(parents=True, exist_ok=True)

        # Loaded IR state
        self._current_ir: Optional[IRFile] = None
        self._ir_data: Optional[np.ndarray] = None         # float32 time-domain
        self._ir_fft_padded: Optional[np.ndarray] = None   # complex FFT at convolution size
        self._conv_fft_len: int = 0                        # cached FFT size (n for rfft)
        self._ir_len: int = 0

        # Overlap-add tail state
        self._tail: np.ndarray = np.array([], dtype=np.float32)

        # Mix & toggle
        self._enabled: bool = True
        self._wet: float = 1.0
        self._dry: float = 0.0

    # ── Public API ──────────────────────────────────────────────────

    def load_ir(self, ir_path: str | Path) -> bool:
        """Load an IR file from disk.

        Args:
            ir_path: Path to .wav IR file.

        Returns:
            True if successfully loaded.
        """
        path = Path(ir_path)
        if not path.exists() or path.suffix not in (".wav",):
            logger.error("IR file not found or invalid: %s", ir_path)
            return False

        sr, data = wavfile.read(path)

        # Normalize to float32 [-1, 1]
        data = self._normalise_wav(data)

        # Mono — take first channel if multi-channel
        if data.ndim > 1:
            data = data[:, 0]

        num_taps = len(data)
        length_ms = (num_taps / sr) * 1000.0

        if num_taps > _MAX_IR_TAPS:
            logger.warning(
                "IR %s has %d taps (max %d); truncating",
                path.stem, num_taps, _MAX_IR_TAPS,
            )
            data = data[:_MAX_IR_TAPS]
            num_taps = len(data)
            length_ms = (num_taps / sr) * 1000.0

        self._current_ir = IRFile(
            name=path.stem,
            path=str(path),
            sample_rate=sr,
            num_taps=num_taps,
            length_ms=length_ms,
        )
        self._ir_data = data
        self._ir_len = num_taps

        # Invalidate cached FFT — will be recomputed at the right size
        # on first process() call with the actual block size
        self._ir_fft_padded = None
        self._conv_fft_len = 0
        self._tail = np.array([], dtype=np.float32)

        logger.info(
            "Loaded IR: %s (%d taps, %.1fms @ %dHz)",
            path.stem, num_taps, length_ms, sr,
        )
        return True

    def process(self, audio_in: np.ndarray) -> np.ndarray:
        """Apply IR convolution using FFT overlap-add.

        Args:
            audio_in: float32 PCM samples in [-1, 1].

        Returns:
            Convolved audio block (same length as input), float32 in [-1, 1].
        """
        # Bypass if no IR loaded or disabled
        if not self._enabled or self._ir_data is None:
            return audio_in

        block_len = len(audio_in)
        ir_len = self._ir_len

        # FFT size for overlap-add: next power of 2 >= block_len + ir_len - 1
        fft_len = _next_pow2(block_len + ir_len - 1)

        # Recompute padded IR FFT if size changed (only on first call or IR reload)
        if fft_len != self._conv_fft_len:
            # ir_data is guaranteed non-None at this point (guard above)
            ir_data: np.ndarray = self._ir_data
            self._ir_fft_padded = np.fft.rfft(ir_data, n=fft_len)
            self._conv_fft_len = fft_len

        # ir_fft_padded is guaranteed non-None after the lazy-init block above
        ir_fft: np.ndarray = self._ir_fft_padded  # type: ignore[assignment]

        # FFT of input block
        block_fft = np.fft.rfft(audio_in, n=fft_len)

        # Multiply in frequency domain (complex element-wise)
        out_fft = block_fft * ir_fft

        # IFFT back to time domain
        convolved = np.fft.irfft(out_fft, n=fft_len).astype(np.float32)

        # Overlap-add: add previous tail to start of current block
        tail_len = len(self._tail)
        if tail_len > 0:
            convolved[:tail_len] += self._tail

        # Save next tail for subsequent block
        tail_slice = convolved[block_len:block_len + ir_len - 1]
        if len(tail_slice) > 0:
            self._tail = tail_slice.copy()
        else:
            self._tail = np.array([], dtype=np.float32)

        # First block_len samples are the valid convolution output
        wet = convolved[:block_len]

        # Wet/dry mix
        if self._wet == 1.0 and self._dry == 0.0:
            return np.clip(wet, -1.0, 1.0)

        return np.clip(
            wet * self._wet + audio_in * self._dry,
            -1.0, 1.0,
        )

    def set_mix(self, wet: float = 1.0, dry: float = 0.0) -> None:
        """Set wet/dry mix levels.

        Args:
            wet: Wet level (0.0-1.0). Fraction of convolved signal.
            dry: Dry level (0.0-1.0). Fraction of original signal.
        """
        self._wet = max(0.0, min(1.0, wet))
        self._dry = max(0.0, min(1.0, dry))

    def reset_state(self) -> None:
        """Reset overlap-add tail state (e.g. on preset switch)."""
        self._tail = np.array([], dtype=np.float32)

    @property
    def wet(self) -> float:
        return self._wet

    @wet.setter
    def wet(self, value: float) -> None:
        self._wet = max(0.0, min(1.0, value))

    @property
    def dry(self) -> float:
        return self._dry

    @dry.setter
    def dry(self, value: float) -> None:
        self._dry = max(0.0, min(1.0, value))

    # ── Query / listing ─────────────────────────────────────────────

    def get_irs(self) -> list[IRFile]:
        """List all available IR files in the IR directory."""
        irs: list[IRFile] = []
        for f in sorted(self._ir_dir.glob("*.wav")):
            try:
                sr, data = wavfile.read(f)
                num_taps = len(data)
                length_ms = (num_taps / sr) * 1000
                channels = data.ndim if data.ndim > 1 else 1
                irs.append(IRFile(
                    name=f.stem,
                    path=str(f),
                    sample_rate=sr,
                    num_taps=num_taps,
                    length_ms=length_ms,
                    channels=channels,
                ))
            except Exception as e:
                logger.warning("Could not read IR %s: %s", f.name, e)
        return irs

    def unload(self) -> None:
        """Unload the current IR and reset state."""
        self._current_ir = None
        self._ir_data = None
        self._ir_fft_padded = None
        self._conv_fft_len = 0
        self._ir_len = 0
        self._tail = np.array([], dtype=np.float32)

    # ── Properties ──────────────────────────────────────────────────

    @property
    def is_loaded(self) -> bool:
        return self._current_ir is not None

    @property
    def current_ir(self) -> Optional[IRFile]:
        return self._current_ir

    @property
    def enabled(self) -> bool:
        return self._enabled

    @enabled.setter
    def enabled(self, value: bool) -> None:
        self._enabled = bool(value)
        if not self._enabled:
            # Clear tail when disabling — next enable starts clean
            self._tail = np.array([], dtype=np.float32)
            logger.debug("IR convolution disabled, tail cleared")

    def reset_tail(self) -> None:
        """Zero out the overlap-add tail (for preset switches)."""
        self._tail = np.array([], dtype=np.float32)

    # ── Internal helpers ────────────────────────────────────────────

    @staticmethod
    def _normalise_wav(data: np.ndarray) -> np.ndarray:
        """Normalise raw WAV data to float32 [-1, 1]."""
        if data.dtype == np.int16:
            return (data.astype(np.float32) / np.float32(32768.0))
        elif data.dtype == np.int32:
            return (data.astype(np.float32) / np.float32(2147483648.0))
        elif data.dtype == np.uint8:
            return ((data.astype(np.float32) - np.float32(128.0))
                    / np.float32(128.0))
        elif data.dtype == np.float32:
            return data
        else:
            return data.astype(np.float32)