"""Shared audio analysis — FFT spectrum, RMS, beat detection.

Engines provide raw audio chunks; AudioAnalyzer processes them into
AudioAnalysis snapshots consumed by visualization streams.
"""

import math
import time
from dataclasses import dataclass, field
from typing import List, Tuple

import numpy as np

# Number of logarithmic frequency bands for spectrum analysis
NUM_BANDS = 64

# Audio defaults
DEFAULT_SAMPLE_RATE = 44100
DEFAULT_CHUNK_SIZE = 2048  # ~46 ms at 44100 Hz


@dataclass
class AudioAnalysis:
    """Snapshot of audio analysis results.

    Written by the capture thread, read by visualization streams.
    Mono fields contain the mixed-down signal (all channels averaged).
    Per-channel fields (left/right) are populated when the source is stereo+.
    For mono sources, left/right are copies of the mono data.
    """

    timestamp: float = 0.0
    # Mono (mixed) — backward-compatible fields
    rms: float = 0.0
    peak: float = 0.0
    spectrum: np.ndarray = field(default_factory=lambda: np.zeros(NUM_BANDS, dtype=np.float32))
    beat: bool = False
    beat_intensity: float = 0.0
    # Per-channel
    left_rms: float = 0.0
    left_spectrum: np.ndarray = field(default_factory=lambda: np.zeros(NUM_BANDS, dtype=np.float32))
    right_rms: float = 0.0
    right_spectrum: np.ndarray = field(default_factory=lambda: np.zeros(NUM_BANDS, dtype=np.float32))


def _build_log_bands(num_bands: int, fft_size: int, sample_rate: int) -> List[Tuple[int, int]]:
    """Build logarithmically-spaced frequency band boundaries for FFT bins.

    Returns list of (start_bin, end_bin) pairs.
    """
    nyquist = sample_rate / 2
    min_freq = 20.0
    max_freq = min(nyquist, 20000.0)
    log_min = math.log10(min_freq)
    log_max = math.log10(max_freq)

    freqs = np.logspace(log_min, log_max, num_bands + 1)
    bin_width = sample_rate / fft_size

    bands = []
    for i in range(num_bands):
        start_bin = max(1, int(freqs[i] / bin_width))
        end_bin = max(start_bin + 1, int(freqs[i + 1] / bin_width))
        end_bin = min(end_bin, fft_size // 2)
        bands.append((start_bin, end_bin))
    return bands


class AudioAnalyzer:
    """Stateful audio analyzer — call analyze() per raw chunk.

    Maintains smoothing buffers, energy history for beat detection,
    and pre-allocated FFT scratch buffers. Thread-safe only if a single
    thread calls analyze() (the capture thread).
    """

    def __init__(self, sample_rate: int = DEFAULT_SAMPLE_RATE, chunk_size: int = DEFAULT_CHUNK_SIZE):
        self._sample_rate = sample_rate
        self._chunk_size = chunk_size

        # FFT helpers
        self._window = np.hanning(chunk_size).astype(np.float32)
        self._bands = _build_log_bands(NUM_BANDS, chunk_size, sample_rate)

        # Beat detection state
        self._energy_history: np.ndarray = np.zeros(43, dtype=np.float64)  # ~1s at 44100/2048
        self._energy_idx = 0

        # Smoothed spectrum (exponential decay)
        self._smooth_spectrum = np.zeros(NUM_BANDS, dtype=np.float32)
        self._smooth_spectrum_left = np.zeros(NUM_BANDS, dtype=np.float32)
        self._smooth_spectrum_right = np.zeros(NUM_BANDS, dtype=np.float32)
        self._smoothing_alpha = 0.3

        # Pre-allocated scratch buffers
        self._fft_windowed = np.empty(chunk_size, dtype=np.float32)
        self._spectrum_buf = np.zeros(NUM_BANDS, dtype=np.float32)
        self._spectrum_buf_left = np.zeros(NUM_BANDS, dtype=np.float32)
        self._spectrum_buf_right = np.zeros(NUM_BANDS, dtype=np.float32)
        self._sq_buf = np.empty(chunk_size, dtype=np.float32)

        # Double-buffered output spectra — avoids allocating new arrays each
        # analyze() call.  Consumers hold a reference to the "old" buffer while
        # the analyzer writes into the alternate one.
        self._out_spectrum = [np.zeros(NUM_BANDS, dtype=np.float32),
                              np.zeros(NUM_BANDS, dtype=np.float32)]
        self._out_spectrum_left = [np.zeros(NUM_BANDS, dtype=np.float32),
                                   np.zeros(NUM_BANDS, dtype=np.float32)]
        self._out_spectrum_right = [np.zeros(NUM_BANDS, dtype=np.float32),
                                    np.zeros(NUM_BANDS, dtype=np.float32)]
        self._out_idx = 0  # toggles 0/1 each analyze() call

        # Pre-compute band start/end arrays and widths for vectorized binning
        self._band_starts = np.array([s for s, _ in self._bands], dtype=np.intp)
        self._band_ends = np.array([e for _, e in self._bands], dtype=np.intp)
        self._band_widths = (self._band_ends - self._band_starts).astype(np.float32)
        self._band_widths[self._band_widths == 0] = 1.0  # avoid divide-by-zero

        # Pre-allocated channel buffers for stereo
        self._left_buf = np.empty(chunk_size, dtype=np.float32)
        self._right_buf = np.empty(chunk_size, dtype=np.float32)
        self._mono_buf = np.empty(chunk_size, dtype=np.float32)

    @property
    def sample_rate(self) -> int:
        return self._sample_rate

    @sample_rate.setter
    def sample_rate(self, value: int):
        if value != self._sample_rate:
            self._sample_rate = value
            self._bands = _build_log_bands(NUM_BANDS, self._chunk_size, value)

    def analyze(self, raw_data: np.ndarray, channels: int) -> AudioAnalysis:
        """Analyze a raw audio chunk and return an AudioAnalysis snapshot.

        Args:
            raw_data: 1-D float32 array of interleaved samples (length = chunk_size * channels)
            channels: Number of audio channels

        Returns:
            AudioAnalysis with spectrum, RMS, beat, etc.
        """
        alpha = self._smoothing_alpha
        one_minus_alpha = 1.0 - alpha

        # Split channels and mix to mono
        if channels > 1:
            data = raw_data.reshape(-1, channels)
            np.copyto(self._left_buf[:len(data)], data[:, 0])
            right_col = data[:, 1] if channels >= 2 else data[:, 0]
            np.copyto(self._right_buf[:len(data)], right_col)
            np.add(data[:, 0], right_col, out=self._mono_buf[:len(data)])
            self._mono_buf[:len(data)] *= 0.5
            samples = self._mono_buf[:len(data)]
            left_samples = self._left_buf[:len(data)]
            right_samples = self._right_buf[:len(data)]
        else:
            samples = raw_data
            left_samples = samples
            right_samples = samples

        # RMS and peak
        n = len(samples)
        np.multiply(samples, samples, out=self._sq_buf[:n])
        rms = float(np.sqrt(np.mean(self._sq_buf[:n])))
        peak = float(np.max(np.abs(samples)))

        if channels > 1:
            np.multiply(left_samples, left_samples, out=self._sq_buf[:n])
            left_rms = float(np.sqrt(np.mean(self._sq_buf[:n])))
            np.multiply(right_samples, right_samples, out=self._sq_buf[:n])
            right_rms = float(np.sqrt(np.mean(self._sq_buf[:n])))
        else:
            left_rms = rms
            right_rms = rms

        # FFT for mono, left, right
        self._fft_bands(samples, self._spectrum_buf, self._smooth_spectrum,
                        alpha, one_minus_alpha)
        if channels > 1:
            self._fft_bands(left_samples, self._spectrum_buf_left, self._smooth_spectrum_left,
                            alpha, one_minus_alpha)
            self._fft_bands(right_samples, self._spectrum_buf_right, self._smooth_spectrum_right,
                            alpha, one_minus_alpha)
        else:
            np.copyto(self._smooth_spectrum_left, self._smooth_spectrum)
            np.copyto(self._smooth_spectrum_right, self._smooth_spectrum)

        # Beat detection — compare current energy to rolling average (mono)
        np.multiply(samples, samples, out=self._sq_buf[:n])
        energy = float(np.sum(self._sq_buf[:n]))
        self._energy_history[self._energy_idx] = energy
        self._energy_idx = (self._energy_idx + 1) % len(self._energy_history)
        avg_energy = float(np.mean(self._energy_history))

        beat = False
        beat_intensity = 0.0
        if avg_energy > 1e-8:
            ratio = energy / avg_energy
            if ratio > 1.5:
                beat = True
                beat_intensity = min(1.0, (ratio - 1.0) / 2.0)

        # Snapshot spectra into double-buffered output arrays (no allocation)
        idx = self._out_idx
        self._out_idx = 1 - idx
        out_spec = self._out_spectrum[idx]
        out_left = self._out_spectrum_left[idx]
        out_right = self._out_spectrum_right[idx]
        np.copyto(out_spec, self._smooth_spectrum)
        np.copyto(out_left, self._smooth_spectrum_left)
        np.copyto(out_right, self._smooth_spectrum_right)

        return AudioAnalysis(
            timestamp=time.perf_counter(),
            rms=rms,
            peak=peak,
            spectrum=out_spec,
            beat=beat,
            beat_intensity=beat_intensity,
            left_rms=left_rms,
            left_spectrum=out_left,
            right_rms=right_rms,
            right_spectrum=out_right,
        )

    def _fft_bands(self, samps, buf, smooth_buf, alpha, one_minus_alpha):
        """Compute FFT, bin into bands, normalize, and smooth."""
        chunk_size = self._chunk_size
        chunk = samps[:chunk_size]
        if len(chunk) < chunk_size:
            chunk = np.pad(chunk, (0, chunk_size - len(chunk)))
        np.multiply(chunk, self._window, out=self._fft_windowed)
        fft_mag = np.abs(np.fft.rfft(self._fft_windowed))
        fft_mag *= (1.0 / chunk_size)
        fft_len = len(fft_mag)
        # Vectorized band binning using cumulative sum
        valid = (self._band_starts < fft_len) & (self._band_ends <= fft_len) & (self._band_ends > 0)
        buf[:] = 0.0
        if valid.any():
            cumsum = np.cumsum(fft_mag)
            band_sums = cumsum[self._band_ends[valid] - 1] - np.where(
                self._band_starts[valid] > 0, cumsum[self._band_starts[valid] - 1], 0.0
            )
            buf[valid] = band_sums / self._band_widths[valid]
        spec_max = float(np.max(buf))
        if spec_max > 1e-6:
            buf *= (1.0 / spec_max)
        smooth_buf *= one_minus_alpha
        smooth_buf += alpha * buf
