"""Audio segmenter with VAD-driven state machine.

Manages speech segment boundaries using Voice Activity Detection.
"""

from __future__ import annotations

from collections import deque
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import TYPE_CHECKING

import numpy as np
from numpy.typing import NDArray

from noteflow.config.constants import DEFAULT_SAMPLE_RATE

if TYPE_CHECKING:
    from collections.abc import Iterator


class SegmenterState(Enum):
    """Segmenter state machine states."""

    IDLE = auto()  # Waiting for speech
    SPEECH = auto()  # Speech detected, accumulating audio
    TRAILING = auto()  # Speech ended, collecting trailing audio


@dataclass
class SegmenterConfig:
    """Configuration for segmenter behavior."""

    # Minimum speech duration to consider valid (seconds)
    min_speech_duration: float = 0.3
    # Maximum segment duration before forced split (seconds)
    max_segment_duration: float = 30.0
    # Trailing silence to include after speech ends (seconds)
    trailing_silence: float = 0.5
    # Leading audio to include before speech starts (seconds)
    leading_buffer: float = 0.2
    # Sample rate for audio processing
    sample_rate: int = DEFAULT_SAMPLE_RATE


@dataclass
class AudioSegment:
    """A completed audio segment ready for transcription."""

    audio: NDArray[np.float32]
    start_time: float
    end_time: float

    @property
    def duration(self) -> float:
        """Segment duration in seconds."""
        return self.end_time - self.start_time


@dataclass
class Segmenter:
    """VAD-driven audio segmenter with state machine.

    Accumulates audio during speech and emits complete segments
    when speech ends or max duration is reached.
    """

    config: SegmenterConfig = field(default_factory=SegmenterConfig)

    # State machine
    _state: SegmenterState = field(default=SegmenterState.IDLE, init=False)

    # Timing tracking
    _stream_time: float = field(default=0.0, init=False)
    _speech_start_time: float = field(default=0.0, init=False)
    _leading_duration: float = field(default=0.0, init=False)

    # Audio buffers with cached sample counts for O(1) length lookups
    # Using deque for _leading_buffer enables O(1) popleft() vs O(n) pop(0)
    _leading_buffer: deque[NDArray[np.float32]] = field(default_factory=deque, init=False)
    _leading_buffer_samples: int = field(default=0, init=False)
    _speech_buffer: list[NDArray[np.float32]] = field(default_factory=list, init=False)
    _speech_buffer_samples: int = field(default=0, init=False)
    _trailing_buffer: list[NDArray[np.float32]] = field(default_factory=list, init=False)
    _trailing_duration: float = field(default=0.0, init=False)

    @property
    def state(self) -> SegmenterState:
        """Get current segmenter state."""
        return self._state

    def reset(self) -> None:
        """Reset segmenter to initial state."""
        self._state = SegmenterState.IDLE
        self._stream_time = 0.0
        self._speech_start_time = 0.0
        self._leading_duration = 0.0
        self._leading_buffer.clear()
        self._leading_buffer_samples = 0
        self._speech_buffer.clear()
        self._speech_buffer_samples = 0
        self._trailing_buffer.clear()
        self._trailing_duration = 0.0

    def process_audio(
        self,
        audio: NDArray[np.float32],
        is_speech: bool,
    ) -> Iterator[AudioSegment]:
        """Process audio chunk with VAD decision.

        Args:
            audio: Audio samples (float32, mono).
            is_speech: VAD decision for this chunk.

        Yields:
            Complete AudioSegment when speech ends or max duration reached.
        """
        chunk_duration = len(audio) / self.config.sample_rate
        chunk_start = self._stream_time
        self._stream_time += chunk_duration

        if self._state == SegmenterState.IDLE:
            yield from self._handle_idle(audio, is_speech, chunk_start)
        elif self._state == SegmenterState.SPEECH:
            yield from self._handle_speech(audio, is_speech, chunk_start, chunk_duration)
        elif self._state == SegmenterState.TRAILING:
            yield from self._handle_trailing(audio, is_speech, chunk_start, chunk_duration)

    def flush(self) -> AudioSegment | None:
        """Flush any pending audio as a segment.

        Call when stream ends to get final segment.

        Returns:
            Remaining audio segment if valid, None otherwise.
        """
        if self._state in (SegmenterState.SPEECH, SegmenterState.TRAILING):
            segment = self._emit_segment()
            self._state = SegmenterState.IDLE
            return segment
        return None

    def _handle_idle(
        self,
        audio: NDArray[np.float32],
        is_speech: bool,
        chunk_start: float,
    ) -> Iterator[AudioSegment]:
        """Handle audio in IDLE state."""
        if is_speech:
            # Speech started - transition to SPEECH state
            self._state = SegmenterState.SPEECH
            self._speech_start_time = chunk_start

            # Capture how much pre-speech audio we are including (O(1) lookup).
            self._leading_duration = self._leading_buffer_samples / self.config.sample_rate

            # Include leading buffer (pre-speech audio)
            self._speech_buffer = list(self._leading_buffer)
            self._speech_buffer_samples = self._leading_buffer_samples + len(audio)
            self._speech_buffer.append(audio)
            self._leading_buffer.clear()
            self._leading_buffer_samples = 0
        else:
            # Still idle - maintain leading buffer
            self._update_leading_buffer(audio)

        yield from ()  # No segments emitted in IDLE

    def _handle_speech(
        self,
        audio: NDArray[np.float32],
        is_speech: bool,
        chunk_start: float,
        chunk_duration: float,
    ) -> Iterator[AudioSegment]:
        """Handle audio in SPEECH state."""
        if is_speech:
            self._speech_buffer.append(audio)
            self._speech_buffer_samples += len(audio)
            current_duration = self._stream_time - self._speech_start_time

            # Check max duration limit
            if current_duration >= self.config.max_segment_duration:
                segment = self._emit_segment()
                if segment is not None:
                    yield segment
                # Start a fresh segment at the end of this chunk
                self._speech_start_time = self._stream_time
                self._leading_duration = 0.0
                self._speech_buffer = []
                self._speech_buffer_samples = 0
        else:
            # Speech ended - transition to TRAILING
            # Start trailing buffer with this silent chunk
            self._state = SegmenterState.TRAILING
            self._trailing_buffer = [audio]
            self._trailing_duration = chunk_duration

            # Check if already past trailing threshold
            if self._trailing_duration >= self.config.trailing_silence:
                segment = self._emit_segment()
                if segment is not None:
                    yield segment
                self._state = SegmenterState.IDLE

    def _handle_trailing(
        self,
        audio: NDArray[np.float32],
        is_speech: bool,
        chunk_start: float,
        chunk_duration: float,
    ) -> Iterator[AudioSegment]:
        """Handle audio in TRAILING state."""
        if is_speech:
            # Speech resumed - merge trailing back and continue
            self._speech_buffer.extend(self._trailing_buffer)
            self._speech_buffer.append(audio)
            self._trailing_buffer.clear()
            self._trailing_duration = 0.0
            self._state = SegmenterState.SPEECH
        else:
            # Still silence - accumulate trailing
            self._trailing_buffer.append(audio)
            self._trailing_duration += chunk_duration

            if self._trailing_duration >= self.config.trailing_silence:
                # Enough trailing silence - emit segment
                segment = self._emit_segment()
                if segment is not None:
                    yield segment
                self._state = SegmenterState.IDLE

    def _update_leading_buffer(self, audio: NDArray[np.float32]) -> None:
        """Maintain rolling leading buffer with O(1) sample counting."""
        self._leading_buffer.append(audio)
        self._leading_buffer_samples += len(audio)

        # Calculate total buffer duration using cached sample count
        total_duration = self._leading_buffer_samples / self.config.sample_rate

        # Trim to configured leading buffer size (O(1) with deque.popleft)
        while total_duration > self.config.leading_buffer and self._leading_buffer:
            removed = self._leading_buffer.popleft()
            self._leading_buffer_samples -= len(removed)
            total_duration = self._leading_buffer_samples / self.config.sample_rate

    def _emit_segment(self) -> AudioSegment | None:
        """Create and emit completed segment."""
        # Combine speech + trailing audio
        all_audio = self._speech_buffer + self._trailing_buffer

        # Calculate actual start time (account for leading buffer)
        actual_start = max(0.0, self._speech_start_time - self._leading_duration)

        # Concatenate audio
        audio = np.concatenate(all_audio) if all_audio else np.array([], dtype=np.float32)

        # If we only have silence/trailing audio, don't emit a segment.
        if not self._speech_buffer:
            self._speech_buffer_samples = 0
            self._trailing_buffer.clear()
            self._trailing_duration = 0.0
            self._leading_duration = 0.0
            return None

        # Check minimum speech duration using cached sample count (O(1))
        speech_duration = self._speech_buffer_samples / self.config.sample_rate
        if speech_duration < self.config.min_speech_duration:
            self._speech_buffer.clear()
            self._speech_buffer_samples = 0
            self._trailing_buffer.clear()
            self._trailing_duration = 0.0
            self._leading_duration = 0.0
            return None

        segment = AudioSegment(
            audio=audio,
            start_time=actual_start,
            end_time=self._stream_time,
        )

        # Clear buffers and reset cached counts
        self._speech_buffer.clear()
        self._speech_buffer_samples = 0
        self._trailing_buffer.clear()
        self._trailing_duration = 0.0
        self._leading_duration = 0.0

        return segment
