from __future__ import annotations

import asyncio
import struct

from fastapi import WebSocketDisconnect

from app.utils.background import BridgeOutputStats, ws_out_loop


class _ClosingWebSocket:
    def __init__(self, exc: BaseException) -> None:
        self.exc = exc
        self.sent_frames = 0

    async def send_bytes(self, frame: bytes) -> None:
        self.sent_frames += 1
        raise self.exc


class _CollectThenDisconnectWebSocket:
    def __init__(self, disconnect_after: int) -> None:
        self.disconnect_after = disconnect_after
        self.sent_frames: list[bytes] = []

    async def send_bytes(self, frame: bytes) -> None:
        self.sent_frames.append(frame)
        if len(self.sent_frames) >= self.disconnect_after:
            raise WebSocketDisconnect(code=1000)

class _CaptureThenDisconnectWebSocket:
    def __init__(self) -> None:
        self.sent_frames: list[bytes] = []

    async def send_bytes(self, frame: bytes) -> None:
        self.sent_frames.append(frame)
        raise WebSocketDisconnect(code=1000)


async def _capture_one_agent_frame(frame: bytes, **ws_out_kwargs) -> _CaptureThenDisconnectWebSocket:
    ws = _CaptureThenDisconnectWebSocket()
    agent_q: asyncio.Queue[bytes] = asyncio.Queue()
    await agent_q.put(frame)

    await asyncio.wait_for(
        ws_out_loop(
            ws,
            agent_q,
            frame_ms=20,
            bytes_per_frame=len(frame),
            **ws_out_kwargs,
        ),
        timeout=0.2,
    )

    return ws


def test_ws_out_loop_exits_when_close_was_already_sent() -> None:
    async def _run() -> None:
        ws = _ClosingWebSocket(RuntimeError('Cannot call "send" once a close message has been sent.'))

        await asyncio.wait_for(
            ws_out_loop(
                ws,
                asyncio.Queue(),
                frame_ms=20,
                bytes_per_frame=4,
            ),
            timeout=0.2,
        )

        assert ws.sent_frames == 1

    asyncio.run(_run())


def test_ws_out_loop_applies_default_output_gain(monkeypatch) -> None:
    monkeypatch.delenv("WS_OUTPUT_GAIN", raising=False)

    frame = struct.pack("<hh", 1000, -1000)
    ws = asyncio.run(_capture_one_agent_frame(frame))

    assert struct.unpack("<hh", ws.sent_frames[0]) == (1000, -1000)


def test_ws_out_loop_reads_output_gain_from_env(monkeypatch) -> None:
    monkeypatch.setenv("WS_OUTPUT_GAIN", "1.5")

    frame = struct.pack("<hh", 1000, -1000)
    ws = asyncio.run(_capture_one_agent_frame(frame))

    assert struct.unpack("<hh", ws.sent_frames[0]) == (1500, -1500)


def test_ws_out_loop_prefers_call_config_output_gain(monkeypatch) -> None:
    monkeypatch.setenv("WS_OUTPUT_GAIN", "1.5")

    frame = struct.pack("<hh", 1000, -1000)
    ws = asyncio.run(_capture_one_agent_frame(frame, output_gain=1.0))

    assert struct.unpack("<hh", ws.sent_frames[0]) == (1000, -1000)


def test_ws_out_loop_exits_on_websocket_disconnect() -> None:
    async def _run() -> None:
        ws = _ClosingWebSocket(WebSocketDisconnect(code=1000))

        await asyncio.wait_for(
            ws_out_loop(
                ws,
                asyncio.Queue(),
                frame_ms=20,
                bytes_per_frame=4,
            ),
            timeout=0.2,
        )

        assert ws.sent_frames == 1

    asyncio.run(_run())


def test_ws_out_loop_waits_for_speech_before_signaling_first_agent_audio() -> None:
    async def _run() -> None:
        silent_frame = struct.pack("<hh", 0, 0)
        voiced_frame = struct.pack("<hh", 2000, -2000)
        ws = _CollectThenDisconnectWebSocket(disconnect_after=3)
        agent_q: asyncio.Queue[bytes] = asyncio.Queue()
        await agent_q.put(silent_frame)
        await agent_q.put(voiced_frame)
        first_audio = asyncio.Event()
        stats = BridgeOutputStats()

        await asyncio.wait_for(
            ws_out_loop(
                ws,
                agent_q,
                frame_ms=20,
                bytes_per_frame=len(silent_frame),
                first_agent_audio_sent=first_audio,
                output_stats=stats,
            ),
            timeout=0.2,
        )

        assert ws.sent_frames[:2] == [silent_frame, voiced_frame]
        assert first_audio.is_set()
        assert stats.first_agent_audio_sent is True
        assert stats.agent_audio_bursts == 1

    asyncio.run(_run())