from __future__ import annotations

from dataclasses import dataclass
from enum import Enum

import numpy as np


class AudioEventKind(str, Enum):
    SPEECH_STARTED = "speech_started"
    PARTIAL_READY = "partial_ready"
    FINAL_READY = "final_ready"


@dataclass(frozen=True)
class AudioEvent:
    kind: AudioEventKind
    samples: np.ndarray | None = None


class AudioSegmenter:
    def __init__(
        self,
        *,
        sample_rate: int = 16000,
        speech_threshold: float = 0.012,
        silence_ms: int = 700,
        partial_interval_ms: int = 1200,
        max_utterance_seconds: int = 30,
        pre_roll_ms: int = 200,
    ) -> None:
        self.sample_rate = sample_rate
        self.speech_threshold = speech_threshold
        self.silence_samples = sample_rate * silence_ms // 1000
        self.partial_samples = sample_rate * partial_interval_ms // 1000
        self.max_samples = sample_rate * max_utterance_seconds
        self.pre_roll_samples = sample_rate * pre_roll_ms // 1000
        self._pre_roll = np.empty(0, dtype=np.float32)
        self._utterance = np.empty(0, dtype=np.float32)
        self._silence_count = 0
        self._last_partial_size = 0
        self._speaking = False

    @property
    def speaking(self) -> bool:
        return self._speaking

    def feed(self, samples: np.ndarray) -> list[AudioEvent]:
        chunk = np.asarray(samples, dtype=np.float32).reshape(-1)
        if chunk.size == 0:
            return []
        rms = float(np.sqrt(np.mean(np.square(chunk), dtype=np.float64)))
        is_speech = rms >= self.speech_threshold
        events: list[AudioEvent] = []

        if not self._speaking:
            if not is_speech:
                self._append_pre_roll(chunk)
                return events
            self._speaking = True
            self._utterance = np.concatenate((self._pre_roll, chunk))
            self._pre_roll = np.empty(0, dtype=np.float32)
            self._silence_count = 0
            self._last_partial_size = 0
            events.append(AudioEvent(AudioEventKind.SPEECH_STARTED))
        else:
            self._utterance = np.concatenate((self._utterance, chunk))

        self._silence_count = 0 if is_speech else self._silence_count + chunk.size
        utterance_size = self._utterance.size

        if (
            utterance_size - self._last_partial_size >= self.partial_samples
            and self._silence_count < self.silence_samples
        ):
            self._last_partial_size = utterance_size
            events.append(
                AudioEvent(
                    AudioEventKind.PARTIAL_READY,
                    self._utterance.copy(),
                )
            )

        if (
            self._silence_count >= self.silence_samples
            or utterance_size >= self.max_samples
        ):
            final = self._utterance.copy()
            self._reset()
            events.append(AudioEvent(AudioEventKind.FINAL_READY, final))

        return events

    def flush(self) -> AudioEvent | None:
        if not self._speaking or self._utterance.size == 0:
            self._reset()
            return None
        final = self._utterance.copy()
        self._reset()
        return AudioEvent(AudioEventKind.FINAL_READY, final)

    def _append_pre_roll(self, chunk: np.ndarray) -> None:
        self._pre_roll = np.concatenate((self._pre_roll, chunk))
        if self._pre_roll.size > self.pre_roll_samples:
            self._pre_roll = self._pre_roll[-self.pre_roll_samples :]

    def _reset(self) -> None:
        self._pre_roll = np.empty(0, dtype=np.float32)
        self._utterance = np.empty(0, dtype=np.float32)
        self._silence_count = 0
        self._last_partial_size = 0
        self._speaking = False