from __future__ import annotations

import time
import unittest

from owner_voice_pet.barge_in import AsyncBargeInMonitor, BargeInSpeakerGate
from owner_voice_pet.models import AudioFrame, AudioSegment
from owner_voice_pet.stt import MetadataSttProvider
from owner_voice_pet.transport import MemoryAudioTransport
from owner_voice_pet.vad import EnergyVadProvider


def speech_frame(idx: int, speaker_id: str, partial: str = "等一下") -> AudioFrame:
    return AudioFrame(
        b"\xff\x7f" * 320,
        16000,
        1,
        idx * 20,
        idx,
        {
            "duration_ms": 20,
            "speech": True,
            "speaker_id": speaker_id,
            "partial_transcript": partial,
        },
    )


def segment_for_speaker(speaker_id: str) -> AudioSegment:
    return AudioSegment(
        b"\xff\x7f" * 640,
        16000,
        1,
        0,
        40,
        {"speaker_id": speaker_id},
    )


class BargeInTests(unittest.TestCase):
    def test_async_monitor_ignores_assistant_playback_voice(self) -> None:
        gate = BargeInSpeakerGate(
            enabled=True,
            user_similarity_threshold=0.62,
            assistant_reject_threshold=0.72,
            min_rms=0.001,
        )
        gate.remember_user_segment(segment_for_speaker("owner"))
        transport = MemoryAudioTransport([speech_frame(1, "assistant"), speech_frame(2, "assistant")])
        transport.start_input()
        monitor = AsyncBargeInMonitor(
            transport=transport,
            vad_provider=EnergyVadProvider(threshold=1),
            realtime_stt=MetadataSttProvider(),
            speaker_gate=gate,
            assistant_profile=gate.assistant_profile(segment_for_speaker("assistant")),
            echo_guard_ms=0,
            min_speech_ms=40,
            listen_interval_ms=1,
        )
        monitor.vad_provider.load()

        monitor.start()
        time.sleep(0.05)
        monitor.stop()

        self.assertFalse(monitor.interrupted)
        self.assertEqual(monitor.pending_frames(), [])

    def test_async_monitor_accepts_user_voice_and_buffers_frames(self) -> None:
        gate = BargeInSpeakerGate(
            enabled=True,
            user_similarity_threshold=0.62,
            assistant_reject_threshold=0.72,
            min_rms=0.001,
        )
        gate.remember_user_segment(segment_for_speaker("owner"))
        transport = MemoryAudioTransport([speech_frame(1, "owner"), speech_frame(2, "owner")])
        transport.start_input()
        monitor = AsyncBargeInMonitor(
            transport=transport,
            vad_provider=EnergyVadProvider(threshold=1),
            realtime_stt=MetadataSttProvider(),
            speaker_gate=gate,
            assistant_profile=gate.assistant_profile(segment_for_speaker("assistant")),
            echo_guard_ms=0,
            min_speech_ms=40,
            listen_interval_ms=1,
        )
        monitor.vad_provider.load()

        monitor.start()
        deadline = time.monotonic() + 1
        while not monitor.interrupted and time.monotonic() < deadline:
            time.sleep(0.005)
        monitor.stop()

        self.assertTrue(monitor.interrupted)
        self.assertEqual([frame.metadata["speaker_id"] for frame in monitor.pending_frames()], ["owner", "owner"])

    def test_user_profile_missing_falls_back_to_non_assistant_with_partial(self) -> None:
        gate = BargeInSpeakerGate(
            enabled=True,
            user_similarity_threshold=0.62,
            assistant_reject_threshold=0.72,
            min_rms=0.001,
        )
        transport = MemoryAudioTransport([speech_frame(1, "guest"), speech_frame(2, "guest")])
        transport.start_input()
        monitor = AsyncBargeInMonitor(
            transport=transport,
            vad_provider=EnergyVadProvider(threshold=1),
            realtime_stt=MetadataSttProvider(),
            speaker_gate=gate,
            assistant_profile=gate.assistant_profile(segment_for_speaker("assistant")),
            echo_guard_ms=0,
            min_speech_ms=40,
            listen_interval_ms=1,
        )
        monitor.vad_provider.load()

        monitor.start()
        deadline = time.monotonic() + 1
        while not monitor.interrupted and time.monotonic() < deadline:
            time.sleep(0.005)
        monitor.stop()

        self.assertTrue(monitor.interrupted)


if __name__ == "__main__":
    unittest.main()