"""Stateful music feature extractor — BPM, beat phase, energy, drop detection.

Consumes AudioAnalysis snapshots and produces higher-level MusicFeatures
for semantic music-reactive visualizations.  Pure numpy, no external deps.
"""

import time
from dataclasses import dataclass

import numpy as np

# Frequency band boundaries (indices into 64-band log spectrum)
_BASS_END = 10  # bins 0-9  ≈ 20-250 Hz
_MID_END = 40  # bins 10-39 ≈ 250-4000 Hz
# bins 40-63 ≈ 4000-20000 Hz (treble)

_MAX_BEAT_HISTORY = 30
_BPM_MIN = 40.0
_BPM_MAX = 220.0

# Drop detection timing
_BUILDUP_MIN_DURATION = 1.0  # seconds of rising energy before "buildup"
_DROP_THRESHOLD = 0.5  # energy must drop by this fraction
_DROP_RECOVERY_TIME = 0.5  # seconds in "drop" state before "recovery"
_RECOVERY_DURATION = 1.0  # seconds in "recovery" before returning to idle
_MIN_ENERGY_FOR_BUILDUP = 0.15  # minimum energy to enter buildup


@dataclass(frozen=True)
class MusicFeatures:
    """Snapshot of semantic music features for one frame."""

    bpm: float = 0.0
    beat: bool = False
    beat_intensity: float = 0.0
    beat_phase: float = 0.0  # 0.0-1.0 position in current beat cycle
    energy: float = 0.0  # smoothed RMS 0.0-1.0
    energy_delta: float = 0.0  # rate of change
    bass_energy: float = 0.0
    mid_energy: float = 0.0
    treble_energy: float = 0.0
    drop_state: str = "idle"  # idle | buildup | drop | recovery
    drop_intensity: float = 0.0


class MusicAnalyzer:
    """Extracts semantic music features from AudioAnalysis snapshots.

    Owned by a single stream thread — not thread-safe, no locking needed.
    """

    def __init__(self, sensitivity: float = 1.0):
        self._sensitivity = sensitivity

        # Beat history (timestamps of detected beats)
        self._beat_times: list[float] = []
        self._bpm: float = 0.0
        self._last_beat_time: float = 0.0

        # Energy tracking
        self._smoothed_energy: float = 0.0
        self._prev_energy: float = 0.0
        self._energy_attack = 0.15  # fast attack
        self._energy_release = 0.05  # slow release

        # Band energy smoothing
        self._bass: float = 0.0
        self._mid: float = 0.0
        self._treble: float = 0.0
        self._band_smoothing = 0.25

        # Drop detection state machine
        self._drop_state = "idle"
        self._drop_state_time: float = 0.0  # when current state started
        self._buildup_start_energy: float = 0.0
        self._peak_energy: float = 0.0
        self._drop_intensity: float = 0.0

    @property
    def sensitivity(self) -> float:
        return self._sensitivity

    @sensitivity.setter
    def sensitivity(self, value: float) -> None:
        self._sensitivity = max(0.1, value)

    def update(self, analysis) -> MusicFeatures:
        """Process one AudioAnalysis frame and return MusicFeatures."""
        now = time.perf_counter()
        spectrum = analysis.spectrum
        rms = analysis.rms * self._sensitivity
        beat = analysis.beat
        beat_intensity = analysis.beat_intensity

        # ── Energy envelope (asymmetric smoothing) ──
        alpha = self._energy_attack if rms > self._smoothed_energy else self._energy_release
        self._smoothed_energy += alpha * (rms - self._smoothed_energy)
        energy = min(1.0, self._smoothed_energy)
        energy_delta = energy - self._prev_energy
        self._prev_energy = energy

        # ── Frequency band energy ──
        bass_raw = (
            float(np.mean(spectrum[:_BASS_END])) * self._sensitivity
            if len(spectrum) > _BASS_END
            else 0.0
        )
        mid_raw = (
            float(np.mean(spectrum[_BASS_END:_MID_END])) * self._sensitivity
            if len(spectrum) > _MID_END
            else 0.0
        )
        treble_raw = (
            float(np.mean(spectrum[_MID_END:])) * self._sensitivity
            if len(spectrum) > _MID_END
            else 0.0
        )

        a = self._band_smoothing
        self._bass += a * (min(1.0, bass_raw) - self._bass)
        self._mid += a * (min(1.0, mid_raw) - self._mid)
        self._treble += a * (min(1.0, treble_raw) - self._treble)

        # ── BPM estimation ──
        if beat:
            self._beat_times.append(now)
            self._last_beat_time = now
            # Keep only recent beats
            if len(self._beat_times) > _MAX_BEAT_HISTORY:
                self._beat_times = self._beat_times[-_MAX_BEAT_HISTORY:]

        self._bpm = self._estimate_bpm()

        # ── Beat phase (0-1 position in current beat cycle) ──
        beat_phase = self._compute_beat_phase(now)

        # ── Drop detection ──
        self._update_drop_state(now, energy, energy_delta)

        return MusicFeatures(
            bpm=self._bpm,
            beat=beat,
            beat_intensity=beat_intensity,
            beat_phase=beat_phase,
            energy=energy,
            energy_delta=energy_delta,
            bass_energy=self._bass,
            mid_energy=self._mid,
            treble_energy=self._treble,
            drop_state=self._drop_state,
            drop_intensity=self._drop_intensity,
        )

    def _estimate_bpm(self) -> float:
        """Estimate BPM from recent beat timestamps using median IBI."""
        if len(self._beat_times) < 3:
            return self._bpm  # keep previous estimate

        # Compute inter-beat intervals
        times = self._beat_times
        ibis = [times[i] - times[i - 1] for i in range(1, len(times))]
        # Filter out unreasonable IBIs
        ibis = [ibi for ibi in ibis if 60.0 / _BPM_MAX <= ibi <= 60.0 / _BPM_MIN]
        if not ibis:
            return self._bpm

        median_ibi = sorted(ibis)[len(ibis) // 2]
        raw_bpm = 60.0 / median_ibi

        # Exponential smoothing to avoid jitter
        if self._bpm > 0:
            return self._bpm * 0.85 + raw_bpm * 0.15
        return raw_bpm

    def _compute_beat_phase(self, now: float) -> float:
        """Compute 0-1 phase within the current beat cycle."""
        if self._bpm <= 0 or self._last_beat_time <= 0:
            return 0.0

        beat_period = 60.0 / self._bpm
        elapsed = now - self._last_beat_time
        phase = (elapsed % beat_period) / beat_period
        return min(1.0, max(0.0, phase))

    def _update_drop_state(self, now: float, energy: float, energy_delta: float) -> None:
        """State machine: idle → buildup → drop → recovery → idle."""
        state_duration = now - self._drop_state_time

        if self._drop_state == "idle":
            # Enter buildup when energy is rising and above minimum
            if energy > _MIN_ENERGY_FOR_BUILDUP and energy_delta > 0.001:
                self._drop_state = "buildup"
                self._drop_state_time = now
                self._buildup_start_energy = energy
                self._peak_energy = energy

        elif self._drop_state == "buildup":
            self._peak_energy = max(self._peak_energy, energy)
            # Transition to drop on sudden energy decrease
            if self._peak_energy > 0 and energy < self._peak_energy * (1 - _DROP_THRESHOLD):
                self._drop_state = "drop"
                self._drop_state_time = now
                self._drop_intensity = min(
                    1.0, (self._peak_energy - energy) / max(self._peak_energy, 0.01)
                )
            # Cancel buildup if energy falls below minimum or state too long
            elif energy < _MIN_ENERGY_FOR_BUILDUP or state_duration > 8.0:
                self._drop_state = "idle"
                self._drop_state_time = now

        elif self._drop_state == "drop":
            if state_duration >= _DROP_RECOVERY_TIME:
                self._drop_state = "recovery"
                self._drop_state_time = now

        elif self._drop_state == "recovery":
            # Fade drop intensity during recovery
            progress = min(1.0, state_duration / _RECOVERY_DURATION)
            self._drop_intensity = self._drop_intensity * (1.0 - progress)
            if state_duration >= _RECOVERY_DURATION:
                self._drop_state = "idle"
                self._drop_state_time = now
                self._drop_intensity = 0.0
