"""speech_utils.py — Pure helpers for wake word, VAD, and STT.

No ROS2 dependencies. All functions operate on raw numpy float32 arrays
at 16 kHz mono (the standard for Whisper / Silero VAD).

Tested by test/test_speech_utils.py (no GPU required).
"""

from __future__ import annotations

import math
import struct
from typing import Generator, List, Tuple

# ── Constants ──────────────────────────────────────────────────────────────────
SAMPLE_RATE = 16000          # Hz — standard for Whisper / Silero
CHUNK_MS = 30                # VAD frame size
CHUNK_SAMPLES = int(SAMPLE_RATE * CHUNK_MS / 1000)   # 480 samples
INT16_MAX = 32768.0


# ── Audio helpers ──────────────────────────────────────────────────────────────

def pcm16_to_float32(data: bytes) -> list:
    """Convert raw PCM16 LE bytes → float32 list in [-1.0, 1.0]."""
    n = len(data) // 2
    samples = struct.unpack(f"<{n}h", data[:n * 2])
    return [s / INT16_MAX for s in samples]


def float32_to_pcm16(samples: list) -> bytes:
    """Convert float32 list → PCM16 LE bytes (clipped to ±1.0)."""
    clipped = [max(-1.0, min(1.0, s)) for s in samples]
    ints = [max(-32768, min(32767, int(s * INT16_MAX))) for s in clipped]
    return struct.pack(f"<{len(ints)}h", *ints)


def rms_db(samples: list) -> float:
    """Compute RMS energy in dBFS. Returns -96.0 for silence."""
    if not samples:
        return -96.0
    mean_sq = sum(s * s for s in samples) / len(samples)
    rms = math.sqrt(mean_sq) if mean_sq > 0.0 else 1e-10
    return 20.0 * math.log10(max(rms, 1e-10))


def chunk_audio(samples: list, chunk_size: int = CHUNK_SAMPLES) -> Generator[list, None, None]:
    """Yield fixed-size chunks from a flat sample list."""
    for i in range(0, len(samples) - chunk_size + 1, chunk_size):
        yield samples[i:i + chunk_size]


# ── Simple energy-based VAD (fallback when Silero unavailable) ────────────────

class EnergyVad:
    """Threshold-based energy VAD.

    Useful as a lightweight fallback / unit-testable reference.
    """

    def __init__(
        self,
        threshold_db: float = -35.0,
        onset_frames: int = 2,
        offset_frames: int = 8,
    ) -> None:
        self.threshold_db = threshold_db
        self.onset_frames = onset_frames
        self.offset_frames = offset_frames
        self._above_count = 0
        self._below_count = 0
        self._active = False

    def process(self, chunk: list) -> bool:
        """Process one chunk (CHUNK_SAMPLES long). Returns True if speech active."""
        db = rms_db(chunk)
        if db >= self.threshold_db:
            self._above_count += 1
            self._below_count = 0
            if self._above_count >= self.onset_frames:
                self._active = True
        else:
            self._below_count += 1
            self._above_count = 0
            if self._below_count >= self.offset_frames:
                self._active = False
        return self._active

    def reset(self) -> None:
        self._above_count = 0
        self._below_count = 0
        self._active = False

    @property
    def is_active(self) -> bool:
        return self._active


# ── Utterance segmenter ────────────────────────────────────────────────────────

class UtteranceSegmenter:
    """Accumulates audio chunks into complete utterances using VAD.

    Yields (samples, duration_s) when an utterance ends.
    Pre-roll keeps a buffer before speech onset (catches first phoneme).
    """

    def __init__(
        self,
        vad: EnergyVad,
        pre_roll_frames: int = 5,
        max_duration_s: float = 15.0,
        sample_rate: int = SAMPLE_RATE,
        chunk_samples: int = CHUNK_SAMPLES,
    ) -> None:
        self._vad = vad
        self._pre_roll = pre_roll_frames
        self._max_frames = int(max_duration_s * sample_rate / chunk_samples)
        self._sample_rate = sample_rate
        self._chunk_samples = chunk_samples
        self._pre_buf: list = []   # ring buffer for pre-roll
        self._utt_buf: list = []   # growing utterance buffer
        self._in_utt = False
        self._silence_after = 0

    def push(self, chunk: list) -> list:
        """Push one chunk, return list of completed (samples, duration_s) tuples."""
        speech = self._vad.process(chunk)
        results = []

        if not self._in_utt:
            self._pre_buf.append(chunk)
            if len(self._pre_buf) > self._pre_roll:
                self._pre_buf.pop(0)

            if speech:
                self._in_utt = True
                self._utt_buf = [s for c in self._pre_buf for s in c]
                self._utt_buf.extend(chunk)
                self._silence_after = 0
        else:
            self._utt_buf.extend(chunk)
            if not speech:
                self._silence_after += 1
                if (self._silence_after >= self._vad.offset_frames or
                        len(self._utt_buf) >= self._max_frames * self._chunk_samples):
                    dur = len(self._utt_buf) / self._sample_rate
                    results.append((list(self._utt_buf), dur))
                    self._utt_buf = []
                    self._in_utt = False
                    self._silence_after = 0
            else:
                self._silence_after = 0

        return results


# ── Speaker embedding distance ─────────────────────────────────────────────────

def cosine_similarity(a: list, b: list) -> float:
    """Cosine similarity between two embedding vectors."""
    if len(a) != len(b) or not a:
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    if norm_a < 1e-10 or norm_b < 1e-10:
        return 0.0
    return dot / (norm_a * norm_b)


def identify_speaker(
    embedding: list,
    known_speakers: dict,
    threshold: float = 0.65,
) -> str:
    """Return speaker_id from known_speakers dict or 'unknown'.

    known_speakers: {speaker_id: embedding_list}
    """
    best_id = "unknown"
    best_sim = threshold
    for spk_id, spk_emb in known_speakers.items():
        sim = cosine_similarity(embedding, spk_emb)
        if sim > best_sim:
            best_sim = sim
            best_id = spk_id
    return best_id
