"""Audio-reactive LED color strip stream.

Implements AudioColorStripStream which produces LED color arrays driven by
real-time audio analysis (spectrum, beat detection, RMS levels).

Three visualization modes:
  spectrum    —  FFT frequency bars mapped across LEDs with palette coloring
  beat_pulse  —  full-strip flash on beat detection with exponential decay
  vu_meter    —  volume level fills LEDs like a progress bar
"""

import threading
import time
from typing import Optional

import numpy as np

from wled_controller.core.audio.analysis import NUM_BANDS
from wled_controller.core.audio.audio_capture import AudioCaptureManager
from wled_controller.core.audio.band_filter import apply_band_filter, compute_band_mask
from wled_controller.core.processing.color_strip_stream import ColorStripStream
from wled_controller.core.processing.effect_stream import _build_palette_lut
from wled_controller.utils import get_logger
from wled_controller.utils.timer import high_resolution_timer

logger = get_logger(__name__)


class AudioColorStripStream(ColorStripStream):
    """Color strip stream driven by audio analysis.

    Dispatches to one of three render methods based on visualization_mode:
    spectrum, beat_pulse, vu_meter.

    Polls AudioCaptureStream for latest analysis data and renders to LED array.
    Uses the same lifecycle pattern as EffectColorStripStream: background
    thread, double-buffered output, configure() for auto-sizing.
    """

    def __init__(self, source, audio_capture_manager: AudioCaptureManager, audio_source_store=None, audio_template_store=None):
        self._audio_capture_manager = audio_capture_manager
        self._audio_source_store = audio_source_store
        self._audio_template_store = audio_template_store
        self._audio_stream = None  # acquired on start

        self._colors_lock = threading.Lock()
        self._running = False
        self._thread: Optional[threading.Thread] = None
        self._fps = 30
        self._frame_time = 1.0 / 30

        # Per-frame timing (read by WledTargetProcessor via get_last_timing())
        self._last_timing: dict = {}

        # Beat pulse persistent state
        self._pulse_brightness = 0.0

        # Smoothed spectrum for temporal smoothing between frames
        self._prev_spectrum: Optional[np.ndarray] = None
        self._prev_rms = 0.0

        self._gradient_store = None  # injected by stream manager
        self._update_from_source(source)

    def set_gradient_store(self, gradient_store) -> None:
        """Inject gradient store for palette resolution."""
        self._gradient_store = gradient_store
        self._resolve_palette_lut()

    def _resolve_palette_lut(self) -> None:
        """Build palette LUT from gradient_id or legacy palette name."""
        gradient_id = self._gradient_id
        if gradient_id and self._gradient_store:
            stops = self._gradient_store.resolve_stops(gradient_id)
            if stops:
                custom = [[s["position"], *s["color"]] for s in stops]
                self._palette_lut = _build_palette_lut("custom", custom)
                return
        self._palette_lut = _build_palette_lut(self._palette_name)

    def _update_from_source(self, source) -> None:
        self._visualization_mode = getattr(source, "visualization_mode", "spectrum")
        self._sensitivity = float(getattr(source, "sensitivity", 1.0))
        self._smoothing = float(getattr(source, "smoothing", 0.3))
        self._gradient_id = getattr(source, "gradient_id", None)
        self._palette_name = getattr(source, "palette", "rainbow")
        self._resolve_palette_lut()
        color = getattr(source, "color", None)
        self._color = color if isinstance(color, list) and len(color) == 3 else [0, 255, 0]
        color_peak = getattr(source, "color_peak", None)
        self._color_peak = color_peak if isinstance(color_peak, list) and len(color_peak) == 3 else [255, 0, 0]
        # Pre-computed float arrays for VU meter (avoid per-frame np.array())
        self._color_f = np.array(self._color, dtype=np.float32)
        self._color_peak_f = np.array(self._color_peak, dtype=np.float32)
        self._auto_size = not source.led_count
        self._led_count = source.led_count if source.led_count and source.led_count > 0 else 1
        self._mirror = bool(getattr(source, "mirror", False))

        # Resolve audio device/channel/template via audio_source_id
        audio_source_id = getattr(source, "audio_source_id", "")
        self._audio_source_id = audio_source_id
        self._audio_engine_type = None
        self._audio_engine_config = None
        self._band_mask = None  # precomputed band filter mask (None = full range)
        if audio_source_id and self._audio_source_store:
            try:
                resolved = self._audio_source_store.resolve_audio_source(audio_source_id)
                self._audio_device_index = resolved.device_index
                self._audio_loopback = resolved.is_loopback
                self._audio_channel = resolved.channel
                if resolved.freq_low is not None and resolved.freq_high is not None:
                    self._band_mask = compute_band_mask(resolved.freq_low, resolved.freq_high)
                if resolved.audio_template_id and self._audio_template_store:
                    try:
                        tpl = self._audio_template_store.get_template(resolved.audio_template_id)
                        self._audio_engine_type = tpl.engine_type
                        self._audio_engine_config = tpl.engine_config
                    except ValueError as e:
                        logger.warning("Audio template %s not found, using default engine: %s", resolved.audio_template_id, e)
                        pass
            except ValueError as e:
                logger.warning(f"Failed to resolve audio source {audio_source_id}: {e}")
                self._audio_device_index = -1
                self._audio_loopback = True
                self._audio_channel = "mono"
        else:
            self._audio_device_index = -1
            self._audio_loopback = True
            self._audio_channel = "mono"

        with self._colors_lock:
            self._colors: Optional[np.ndarray] = None

    # ── ColorStripStream interface ──────────────────────────────────

    def configure(self, device_led_count: int) -> None:
        if self._auto_size and device_led_count > 0:
            new_count = max(self._led_count, device_led_count)
            if new_count != self._led_count:
                self._led_count = new_count
                logger.debug(f"AudioColorStripStream auto-sized to {new_count} LEDs")

    @property
    def target_fps(self) -> int:
        return self._fps

    @property
    def led_count(self) -> int:
        return self._led_count

    def set_capture_fps(self, fps: int) -> None:
        self._fps = max(1, min(90, fps))
        self._frame_time = 1.0 / self._fps

    def start(self) -> None:
        if self._running:
            return
        # Acquire shared audio capture stream
        self._audio_stream = self._audio_capture_manager.acquire(
            self._audio_device_index, self._audio_loopback,
            engine_type=self._audio_engine_type,
            engine_config=self._audio_engine_config,
        )
        self._running = True
        self._thread = threading.Thread(
            target=self._animate_loop,
            name=f"css-audio-{self._visualization_mode}",
            daemon=True,
        )
        self._thread.start()
        logger.info(
            f"AudioColorStripStream started (viz={self._visualization_mode}, "
            f"engine={self._audio_engine_type}, "
            f"device={self._audio_device_index}, loopback={self._audio_loopback})"
        )

    def stop(self) -> None:
        self._running = False
        if self._thread:
            self._thread.join(timeout=5.0)
            if self._thread.is_alive():
                logger.warning("AudioColorStripStream thread did not terminate within 5s")
            self._thread = None
        # Release shared audio capture
        if self._audio_stream is not None:
            self._audio_capture_manager.release(
                self._audio_device_index, self._audio_loopback,
                engine_type=self._audio_engine_type,
            )
            self._audio_stream = None
        self._prev_spectrum = None
        logger.info("AudioColorStripStream stopped")

    def get_latest_colors(self) -> Optional[np.ndarray]:
        with self._colors_lock:
            return self._colors

    def get_last_timing(self) -> dict:
        return dict(self._last_timing)

    def update_source(self, source) -> None:
        from wled_controller.storage.color_strip_source import AudioColorStripSource
        if isinstance(source, AudioColorStripSource):
            old_device = self._audio_device_index
            old_loopback = self._audio_loopback
            old_engine_type = self._audio_engine_type
            prev_led_count = self._led_count if self._auto_size else None
            self._update_from_source(source)
            if prev_led_count and self._auto_size:
                self._led_count = prev_led_count

            # If audio device or engine changed, swap capture stream
            needs_swap = (
                self._audio_device_index != old_device
                or self._audio_loopback != old_loopback
                or self._audio_engine_type != old_engine_type
            )
            if self._running and needs_swap:
                self._audio_capture_manager.release(
                    old_device, old_loopback, engine_type=old_engine_type,
                )
                self._audio_stream = self._audio_capture_manager.acquire(
                    self._audio_device_index, self._audio_loopback,
                    engine_type=self._audio_engine_type,
                    engine_config=self._audio_engine_config,
                )
                logger.info(
                    f"AudioColorStripStream swapped audio device: "
                    f"{old_engine_type}:{old_device}:{old_loopback} → "
                    f"{self._audio_engine_type}:{self._audio_device_index}:{self._audio_loopback}"
                )

            logger.info("AudioColorStripStream params updated in-place")

    # ── Main animation loop ─────────────────────────────────────────

    def _animate_loop(self) -> None:
        _pool_n = 0
        _buf_a = _buf_b = None
        _use_a = True
        # Pre-computed interpolation arrays (rebuilt when LED count changes)
        _band_x = None
        _led_x = None
        _led_x_mirror = None
        _full_amp = None
        _vu_gradient = None
        _indices_buf = None

        renderers = {
            "spectrum": self._render_spectrum,
            "beat_pulse": self._render_beat_pulse,
            "vu_meter": self._render_vu_meter,
        }

        try:
            with high_resolution_timer():
                while self._running:
                    loop_start = time.perf_counter()
                    frame_time = self._frame_time
                    try:
                        n = self._led_count

                        # Rebuild scratch buffers and pre-computed arrays when LED count changes
                        if n != _pool_n:
                            _pool_n = n
                            _buf_a = np.zeros((n, 3), dtype=np.uint8)
                            _buf_b = np.zeros((n, 3), dtype=np.uint8)
                            _band_x = np.arange(NUM_BANDS, dtype=np.float32)
                            half = (n + 1) // 2
                            _led_x_mirror = np.linspace(0, NUM_BANDS - 1, half)
                            _led_x = np.linspace(0, NUM_BANDS - 1, n)
                            _full_amp = np.empty(n, dtype=np.float32)
                            _vu_gradient = np.linspace(0, 1, n, dtype=np.float32)
                            _indices_buf = np.empty(n, dtype=np.int32)
                            _f32_rgb = np.empty((n, 3), dtype=np.float32)
                            self._prev_spectrum = None  # reset smoothing on resize

                        # Make pre-computed arrays available to render methods
                        self._band_x = _band_x
                        self._led_x = _led_x
                        self._led_x_mirror = _led_x_mirror
                        self._full_amp = _full_amp
                        self._vu_gradient = _vu_gradient
                        self._indices_buf = _indices_buf
                        self._f32_rgb = _f32_rgb

                        buf = _buf_a if _use_a else _buf_b
                        _use_a = not _use_a

                        # Get latest audio analysis
                        analysis = None
                        if self._audio_stream is not None:
                            analysis = self._audio_stream.get_latest_analysis()

                        render_fn = renderers.get(self._visualization_mode, self._render_spectrum)
                        t_render = time.perf_counter()
                        render_fn(buf, n, analysis)
                        render_ms = (time.perf_counter() - t_render) * 1000

                        with self._colors_lock:
                            self._colors = buf

                        # Pull capture-side timing and combine with render timing
                        capture_timing = self._audio_stream.get_last_timing() if self._audio_stream else {}
                        read_ms = capture_timing.get("read_ms", 0)
                        fft_ms = capture_timing.get("fft_ms", 0)
                        self._last_timing = {
                            "audio_read_ms": read_ms,
                            "audio_fft_ms": fft_ms,
                            "audio_render_ms": render_ms,
                            "total_ms": read_ms + fft_ms + render_ms,
                        }
                    except Exception as e:
                        logger.error(f"AudioColorStripStream render error: {e}")

                    elapsed = time.perf_counter() - loop_start
                    time.sleep(max(frame_time - elapsed, 0.001))
        except Exception as e:
            logger.error(f"Fatal AudioColorStripStream loop error: {e}", exc_info=True)
        finally:
            self._running = False

    # ── Channel selection ─────────────────────────────────────────

    def _pick_channel(self, analysis):
        """Return (spectrum, rms) for the configured audio channel, with band filtering."""
        if self._audio_channel == "left":
            spectrum, rms = analysis.left_spectrum, analysis.left_rms
        elif self._audio_channel == "right":
            spectrum, rms = analysis.right_spectrum, analysis.right_rms
        else:
            spectrum, rms = analysis.spectrum, analysis.rms
        if self._band_mask is not None:
            spectrum, rms = apply_band_filter(spectrum, rms, self._band_mask)
        return spectrum, rms

    # ── Spectrum Analyzer ──────────────────────────────────────────

    def _render_spectrum(self, buf: np.ndarray, n: int, analysis) -> None:
        if analysis is None:
            buf[:] = 0
            return

        spectrum, _ = self._pick_channel(analysis)
        sensitivity = self._sensitivity
        smoothing = self._smoothing
        lut = self._palette_lut
        band_x = self._band_x
        full_amp = self._full_amp
        indices_buf = self._indices_buf

        if self._mirror:
            half = (n + 1) // 2
            amplitudes = np.interp(self._led_x_mirror, band_x, spectrum)
            amplitudes *= sensitivity
            np.clip(amplitudes, 0.0, 1.0, out=amplitudes)
            if self._prev_spectrum is not None and len(self._prev_spectrum) == half:
                amplitudes *= (1.0 - smoothing)
                amplitudes += smoothing * self._prev_spectrum
            self._prev_spectrum = amplitudes.copy()
            # Mirror: center = bass, edges = treble
            full_amp[:half] = amplitudes[::-1]
            full_amp[half:] = amplitudes[: n - half]
        else:
            amplitudes = np.interp(self._led_x, band_x, spectrum)
            amplitudes *= sensitivity
            np.clip(amplitudes, 0.0, 1.0, out=amplitudes)
            if self._prev_spectrum is not None and len(self._prev_spectrum) == n:
                amplitudes *= (1.0 - smoothing)
                amplitudes += smoothing * self._prev_spectrum
            self._prev_spectrum = amplitudes.copy()
            full_amp[:] = amplitudes

        # Map to palette: amplitude → palette index → color
        np.multiply(full_amp, 255, out=full_amp)
        np.clip(full_amp, 0, 255, out=full_amp)
        np.copyto(indices_buf, full_amp, casting='unsafe')
        colors = lut[indices_buf]  # (n, 3) uint8

        # Scale brightness by amplitude — restore full_amp to [0, 1]
        full_amp *= (1.0 / 255.0)
        f32_rgb = self._f32_rgb
        np.copyto(f32_rgb, colors, casting='unsafe')
        f32_rgb *= full_amp[:, np.newaxis]
        np.clip(f32_rgb, 0, 255, out=f32_rgb)
        np.copyto(buf, f32_rgb, casting='unsafe')

    # ── VU Meter ───────────────────────────────────────────────────

    def _render_vu_meter(self, buf: np.ndarray, n: int, analysis) -> None:
        if analysis is None:
            buf[:] = 0
            return

        _, ch_rms = self._pick_channel(analysis)
        rms = ch_rms * self._sensitivity
        rms = self._smoothing * self._prev_rms + (1.0 - self._smoothing) * rms
        self._prev_rms = rms
        rms = min(1.0, rms)

        fill_count = int(rms * n)
        buf[:] = 0

        if fill_count > 0:
            base = self._color_f
            peak = self._color_peak_f
            t = self._vu_gradient[:fill_count]
            for ch in range(3):
                buf[:fill_count, ch] = np.clip(
                    base[ch] + (peak[ch] - base[ch]) * t, 0, 255
                ).astype(np.uint8)

    # ── Beat Pulse ─────────────────────────────────────────────────

    def _render_beat_pulse(self, buf: np.ndarray, n: int, analysis) -> None:
        if analysis is None:
            buf[:] = 0
            return

        if analysis.beat:
            self._pulse_brightness = 1.0
        else:
            decay_rate = 0.05 + 0.15 * (1.0 / max(self._sensitivity, 0.1))
            self._pulse_brightness = max(0.0, self._pulse_brightness - decay_rate)

        brightness = self._pulse_brightness
        if brightness < 0.01:
            buf[:] = 0
            return

        palette_idx = max(0, min(255, int(analysis.beat_intensity * 255)))
        base_color = self._palette_lut[palette_idx]

        # Vectorized fill: scale color by brightness and broadcast to all LEDs
        r, g, b = int(base_color[0] * brightness), int(base_color[1] * brightness), int(base_color[2] * brightness)
        buf[:, 0] = r
        buf[:, 1] = g
        buf[:, 2] = b
