"""VAD-based audio chunker for live meeting transcription.

Receives 16 kHz Int16 mono samples (the canonical Lirelle meeting stream
format — see meeting_recorder.py). Uses webrtcvad to detect natural speech
pauses and emits chunks at those boundaries via a callback so they can be
transcribed incrementally instead of waiting for the meeting to end.

Design choices:

  - We chunk on natural pauses (≥800 ms silence after ≥5 s of speech),
    not on fixed time intervals. Whisper accuracy degrades sharply when
    cuts land mid-sentence; pausing-on-pause keeps each chunk a
    complete-ish utterance.

  - We do NOT chunk silence-only audio. If the speaker hasn't said
    anything yet, samples are dropped from the VAD pipeline (still
    written to the WAV by the drain thread — this class only feeds the
    live transcript path).

  - We keep ~200 ms of trailing silence inside each emitted chunk because
    Whisper's punctuation/intonation prediction needs some terminal
    context. Anything beyond that gets trimmed so the next chunk doesn't
    start with stale silence.

  - Hard cap at 60 s per chunk for the case where someone monologues
    without pause. Force-emits even mid-sentence so the user sees text
    appear during long talking heads.

  - feed() is called from the recorder's drain thread and must be cheap.
    VAD evaluation per 20 ms frame is ~10 μs on Apple Silicon. Emission
    callback should be fast (typically: append to a queue). The actual
    transcription work happens on a separate thread in
    meeting_live_transcribe.

  - Thread-safe — internal lock around the feed/flush mutations so it's
    safe to feed from drain thread + flush from stop thread.
"""

from __future__ import annotations

import threading
from typing import Callable

import numpy as np


# Wire constants. Whisper canonical rate is 16 kHz; that's also the only
# rate Lirelle's meeting WAVs are written at, so no resampling needed here.
SAMPLE_RATE = 16000
FRAME_MS = 20                   # webrtcvad supports 10 / 20 / 30 ms
FRAME_SAMPLES = SAMPLE_RATE * FRAME_MS // 1000   # 320
FRAME_BYTES = FRAME_SAMPLES * 2                   # int16

# Chunking policy.
MIN_CHUNK_DURATION_S = 5.0      # don't emit anything shorter than this
MAX_CHUNK_DURATION_S = 60.0     # force-emit at this length (monologue case)
SILENCE_TRIGGER_MS = 800        # this much trailing silence → emit
KEEP_TAIL_MS = 200              # preserve this much terminal silence in chunk
# Minimum length on flush(): if the partial chunk is shorter than this,
# we drop it rather than burn a transcription call on a fragment.
FLUSH_MIN_DURATION_S = 0.5

# webrtcvad aggressiveness 0..3 — 2 is the sweet spot for meeting audio
# (less false-positive on background noise than 3, less missed speech
# than 0-1).
VAD_AGGRESSIVENESS = 2


ChunkCallback = Callable[[float, np.ndarray], None]


class VADChunker:
    """Per-stream chunker. One instance per speaker (mic vs system).

    Owner provides an `on_chunk(start_offset_s, audio_float32)` callback
    that fires when a chunk is ready. The callback is invoked synchronously
    from inside feed()/flush() — make it fast (queue.put is ideal).

    Sample timeline starts at zero when this instance is created. The
    `start_offset_s` passed to the callback is wall-time-from-start, used
    by the transcription pipeline to align chunk-local timestamps with the
    global meeting timeline.
    """

    def __init__(self, on_chunk: ChunkCallback) -> None:
        self._on_chunk = on_chunk
        # Lazy-import webrtcvad here (not at module top) so a missing
        # transitive (pkg_resources → jaraco.text) only blows up at
        # meeting start — caught by MeetingRecording, which then falls
        # back to recording-only mode. Same pattern as recorder.py's
        # hands-free VAD path.
        import webrtcvad
        self._vad = webrtcvad.Vad(VAD_AGGRESSIVENESS)

        # Input buffer — raw int16 bytes pending VAD frame extraction.
        self._pending = bytearray()
        # Accumulator for the current in-progress chunk (int16 bytes).
        self._chunk = bytearray()
        # Start offset of current chunk within the meeting (seconds).
        self._chunk_start_offset_s = 0.0
        # Consecutive silence frames inside the current chunk.
        self._silence_run = 0
        # True iff we've seen any speech in the current chunk.
        self._in_speech = False
        # Running total of samples processed (for offset math).
        self._total_samples = 0

        self._lock = threading.Lock()

    # ── feed path ──────────────────────────────────────────────────────

    def feed(self, int16_samples: np.ndarray) -> None:
        """Append 16 kHz mono Int16 samples to the chunker. May fire one
        or more on_chunk callbacks before returning."""
        if int16_samples.size == 0:
            return
        with self._lock:
            self._pending.extend(int16_samples.tobytes())
            self._process()

    def _process(self) -> None:
        # Process every whole 20 ms frame available in _pending.
        while len(self._pending) >= FRAME_BYTES:
            frame_bytes = bytes(self._pending[:FRAME_BYTES])
            del self._pending[:FRAME_BYTES]
            frame_start_sample = self._total_samples
            self._total_samples += FRAME_SAMPLES

            try:
                is_speech = self._vad.is_speech(frame_bytes, SAMPLE_RATE)
            except Exception:
                # webrtcvad occasionally raises on degenerate input; treat
                # the frame as silence and move on. This keeps the chunker
                # alive through transient audio glitches.
                is_speech = False

            if is_speech:
                if not self._in_speech:
                    # First speech frame of a new chunk — mark its start.
                    self._chunk_start_offset_s = (
                        frame_start_sample / SAMPLE_RATE
                    )
                    self._in_speech = True
                self._chunk.extend(frame_bytes)
                self._silence_run = 0
            else:
                if self._in_speech:
                    # Trailing silence inside an active chunk — keep
                    # appending; we'll trim later when we emit.
                    self._chunk.extend(frame_bytes)
                    self._silence_run += 1
                    if self._should_emit_on_silence():
                        self._emit_chunk(trim_trailing_silence=True)
                        continue
                # else: leading silence before any speech — discard frame.

            # Hard cap: force-emit if chunk has run too long without a
            # natural pause.
            if self._in_speech and self._chunk_duration_s() >= MAX_CHUNK_DURATION_S:
                self._emit_chunk(trim_trailing_silence=False)

    # ── emission policy ────────────────────────────────────────────────

    def _should_emit_on_silence(self) -> bool:
        silence_ms = self._silence_run * FRAME_MS
        if silence_ms < SILENCE_TRIGGER_MS:
            return False
        # Chunk must include at least MIN_CHUNK_DURATION_S of REAL content
        # (excluding the trailing silence we're about to trim).
        chunk_ms = self._chunk_duration_s() * 1000
        speech_ms = chunk_ms - silence_ms
        return speech_ms >= MIN_CHUNK_DURATION_S * 1000

    def _chunk_duration_s(self) -> float:
        return len(self._chunk) / 2 / SAMPLE_RATE

    def _emit_chunk(self, *, trim_trailing_silence: bool) -> None:
        """Build the chunk array, call on_chunk, reset state."""
        chunk_bytes = bytes(self._chunk)
        if trim_trailing_silence and self._silence_run > 0:
            keep_tail_frames = max(1, KEEP_TAIL_MS // FRAME_MS)
            trim_frames = self._silence_run - keep_tail_frames
            if trim_frames > 0:
                trim_bytes = trim_frames * FRAME_BYTES
                chunk_bytes = chunk_bytes[: len(chunk_bytes) - trim_bytes]

        # Reset BEFORE invoking the callback so a synchronous callback
        # that calls back into this chunker (it shouldn't, but defensively)
        # sees clean state.
        start_offset_s = self._chunk_start_offset_s
        self._chunk = bytearray()
        self._chunk_start_offset_s = 0.0
        self._silence_run = 0
        self._in_speech = False

        if not chunk_bytes:
            return
        int16 = np.frombuffer(chunk_bytes, dtype=np.int16)
        audio_f32 = int16.astype(np.float32) / 32768.0
        try:
            self._on_chunk(start_offset_s, audio_f32)
        except Exception as e:
            print(f"[meeting_chunker] on_chunk callback raised: {e}", flush=True)

    # ── flush path (called once at stop) ──────────────────────────────

    def flush(self) -> None:
        """Force-emit the partial chunk currently being accumulated, if
        it's long enough to be useful. Called by MeetingRecording.stop()
        so the final utterance of the meeting still makes it into the
        transcript even if the speaker didn't pause before stopping."""
        with self._lock:
            if not self._in_speech or self._chunk_duration_s() < FLUSH_MIN_DURATION_S:
                # Reset state regardless so the chunker isn't reusable in
                # a broken state if someone misuses it.
                self._chunk = bytearray()
                self._chunk_start_offset_s = 0.0
                self._silence_run = 0
                self._in_speech = False
                return
            # Trim only if we have enough trailing silence to bother.
            trim = self._silence_run * FRAME_MS >= SILENCE_TRIGGER_MS
            self._emit_chunk(trim_trailing_silence=trim)

    # ── debug helpers ──────────────────────────────────────────────────

    @property
    def total_samples_seen(self) -> int:
        with self._lock:
            return self._total_samples

    @property
    def current_chunk_duration_s(self) -> float:
        with self._lock:
            return self._chunk_duration_s() if self._in_speech else 0.0
