"""test_speech_utils.py — Unit tests for speech_utils.py (no ROS2 / GPU needed)."""

import math
import struct
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from saltybot_social.speech_utils import (
    pcm16_to_float32, float32_to_pcm16,
    rms_db, chunk_audio,
    EnergyVad, UtteranceSegmenter,
    cosine_similarity, identify_speaker,
    CHUNK_SAMPLES, SAMPLE_RATE,
)
import pytest


# ── PCM conversion ────────────────────────────────────────────────────────────

class TestPcmConversion:
    def test_pcm16_to_float32_silence(self):
        data = b"\x00\x00" * 100
        samples = pcm16_to_float32(data)
        assert len(samples) == 100
        assert all(s == 0.0 for s in samples)

    def test_pcm16_to_float32_max_positive(self):
        data = struct.pack("<h", 32767)
        samples = pcm16_to_float32(data)
        assert abs(samples[0] - 32767 / 32768.0) < 1e-5

    def test_float32_to_pcm16_roundtrip(self):
        original = [0.0, 0.5, -0.5, 0.99, -0.99]
        pcm = float32_to_pcm16(original)
        recovered = pcm16_to_float32(pcm)
        for orig, rec in zip(original, recovered):
            assert abs(orig - rec) < 1e-3

    def test_float32_to_pcm16_clips(self):
        # Values > 1.0 should be clipped
        pcm = float32_to_pcm16([2.0, -2.0])
        samples = pcm16_to_float32(pcm)
        assert samples[0] <= 1.0
        assert samples[1] >= -1.0


# ── RMS dB ────────────────────────────────────────────────────────────────────

class TestRmsDb:
    def test_silence_returns_minus96(self):
        assert rms_db([]) == -96.0
        assert rms_db([0.0] * 100) < -60.0

    def test_full_scale_is_near_zero(self):
        samples = [1.0] * 100
        db = rms_db(samples)
        assert abs(db) < 0.1   # 0 dBFS

    def test_half_scale_is_minus6(self):
        samples = [0.5] * 100
        db = rms_db(samples)
        assert abs(db - (-6.02)) < 0.1

    def test_low_level_signal(self):
        samples = [0.001] * 100
        db = rms_db(samples)
        assert db < -40.0


# ── Chunk audio ───────────────────────────────────────────────────────────────

class TestChunkAudio:
    def test_even_chunks(self):
        samples = list(range(480 * 4))
        chunks = list(chunk_audio(samples, 480))
        assert len(chunks) == 4
        assert all(len(c) == 480 for c in chunks)

    def test_remainder_dropped(self):
        samples = list(range(500))
        chunks = list(chunk_audio(samples, 480))
        assert len(chunks) == 1

    def test_empty_input(self):
        assert list(chunk_audio([], 480)) == []


# ── Energy VAD ────────────────────────────────────────────────────────────────

class TestEnergyVad:
    def test_silence_not_active(self):
        vad = EnergyVad(threshold_db=-35.0, onset_frames=2, offset_frames=3)
        silence = [0.0001] * CHUNK_SAMPLES
        for _ in range(5):
            active = vad.process(silence)
        assert not active

    def test_loud_signal_activates(self):
        vad = EnergyVad(threshold_db=-35.0, onset_frames=2, offset_frames=8)
        loud = [0.5] * CHUNK_SAMPLES
        results = [vad.process(loud) for _ in range(3)]
        assert results[-1] is True

    def test_onset_requires_n_frames(self):
        vad = EnergyVad(threshold_db=-35.0, onset_frames=3, offset_frames=8)
        loud = [0.5] * CHUNK_SAMPLES
        assert not vad.process(loud)   # frame 1
        assert not vad.process(loud)   # frame 2
        assert vad.process(loud)       # frame 3 → activates

    def test_offset_deactivates(self):
        vad = EnergyVad(threshold_db=-35.0, onset_frames=1, offset_frames=2)
        loud = [0.5] * CHUNK_SAMPLES
        silence = [0.0001] * CHUNK_SAMPLES
        vad.process(loud)
        assert vad.is_active
        vad.process(silence)
        assert vad.is_active   # offset_frames=2, need 2
        vad.process(silence)
        assert not vad.is_active

    def test_reset(self):
        vad = EnergyVad(threshold_db=-35.0, onset_frames=1, offset_frames=8)
        loud = [0.5] * CHUNK_SAMPLES
        vad.process(loud)
        assert vad.is_active
        vad.reset()
        assert not vad.is_active


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

class TestUtteranceSegmenter:
    def _make_segmenter(self):
        vad = EnergyVad(threshold_db=-35.0, onset_frames=1, offset_frames=2)
        return UtteranceSegmenter(vad, pre_roll_frames=2, max_duration_s=5.0)

    def test_silence_yields_nothing(self):
        seg = self._make_segmenter()
        silence = [0.0001] * CHUNK_SAMPLES
        for _ in range(10):
            results = seg.push(silence)
            assert results == []

    def test_speech_then_silence_yields_utterance(self):
        seg = self._make_segmenter()
        loud = [0.5] * CHUNK_SAMPLES
        silence = [0.0001] * CHUNK_SAMPLES

        # 3 speech frames
        for _ in range(3):
            assert seg.push(loud) == []

        # 2 silence frames to trigger offset
        all_results = []
        for _ in range(3):
            all_results.extend(seg.push(silence))

        assert len(all_results) == 1
        utt_samples, duration = all_results[0]
        assert duration > 0
        assert len(utt_samples) > 0

    def test_preroll_included(self):
        vad = EnergyVad(threshold_db=-35.0, onset_frames=1, offset_frames=2)
        seg = UtteranceSegmenter(vad, pre_roll_frames=3, max_duration_s=5.0)
        silence = [0.0001] * CHUNK_SAMPLES
        loud = [0.5] * CHUNK_SAMPLES

        # Push 3 silence frames (become pre-roll buffer)
        for _ in range(3):
            seg.push(silence)

        # Push 1 speech frame (triggers onset, pre-roll included)
        seg.push(loud)

        # End utterance
        all_results = []
        for _ in range(3):
            all_results.extend(seg.push(silence))

        assert len(all_results) == 1
        # Utterance should include pre-roll + speech + trailing silence
        utt_samples, _ = all_results[0]
        assert len(utt_samples) >= 4 * CHUNK_SAMPLES


# ── Speaker identification ────────────────────────────────────────────────────

class TestSpeakerIdentification:
    def test_cosine_identical(self):
        v = [1.0, 0.0, 0.0]
        assert abs(cosine_similarity(v, v) - 1.0) < 1e-6

    def test_cosine_orthogonal(self):
        a = [1.0, 0.0]
        b = [0.0, 1.0]
        assert abs(cosine_similarity(a, b)) < 1e-6

    def test_cosine_opposite(self):
        a = [1.0, 0.0]
        b = [-1.0, 0.0]
        assert abs(cosine_similarity(a, b) + 1.0) < 1e-6

    def test_cosine_zero_vector(self):
        assert cosine_similarity([0.0, 0.0], [1.0, 0.0]) == 0.0

    def test_cosine_length_mismatch(self):
        assert cosine_similarity([1.0], [1.0, 0.0]) == 0.0

    def test_identify_speaker_match(self):
        known = {
            "alice": [1.0, 0.0, 0.0],
            "bob":   [0.0, 1.0, 0.0],
        }
        result = identify_speaker([1.0, 0.0, 0.0], known, threshold=0.5)
        assert result == "alice"

    def test_identify_speaker_unknown(self):
        known = {"alice": [1.0, 0.0, 0.0]}
        result = identify_speaker([0.0, 1.0, 0.0], known, threshold=0.5)
        assert result == "unknown"

    def test_identify_speaker_empty_db(self):
        result = identify_speaker([1.0, 0.0], {}, threshold=0.5)
        assert result == "unknown"

    def test_identify_best_match(self):
        known = {
            "alice": [1.0, 0.0, 0.0],
            "bob":   [0.8, 0.6, 0.0],  # closer to [1,0,0] than bob is
        }
        # Query close to alice
        result = identify_speaker([0.99, 0.01, 0.0], known, threshold=0.5)
        assert result == "alice"
