"""Unit tests for the WebSocket connection manager."""

import json
from unittest.mock import AsyncMock, MagicMock

import pytest
from agenteval.web.websocket import ConnectionManager


@pytest.fixture()
def manager():
    return ConnectionManager()


async def test_connect_adds_websocket(manager):
    ws = AsyncMock()
    await manager.connect("run-1", ws)
    ws.accept.assert_called_once()
    assert ws in manager._connections["run-1"]


async def test_disconnect_removes_websocket(manager):
    ws = AsyncMock()
    await manager.connect("run-1", ws)
    manager.disconnect("run-1", ws)
    assert "run-1" not in manager._connections


async def test_disconnect_cleans_empty_group(manager):
    ws1 = AsyncMock()
    ws2 = AsyncMock()
    await manager.connect("run-1", ws1)
    await manager.connect("run-1", ws2)
    manager.disconnect("run-1", ws1)
    assert ws2 in manager._connections["run-1"]
    manager.disconnect("run-1", ws2)
    assert "run-1" not in manager._connections


async def test_broadcast_sends_to_all_connections(manager):
    ws1 = AsyncMock()
    ws2 = AsyncMock()
    await manager.connect("run-1", ws1)
    await manager.connect("run-1", ws2)

    await manager.broadcast("run-1", {"event": "test", "data": "value"})

    ws1.send_text.assert_called_once()
    ws2.send_text.assert_called_once()
    msg = json.loads(ws1.send_text.call_args[0][0])
    assert msg["event"] == "test"


async def test_broadcast_removes_dead_connections(manager):
    ws1 = AsyncMock()
    ws1.send_text.side_effect = Exception("connection closed")
    ws2 = AsyncMock()
    await manager.connect("run-1", ws1)
    await manager.connect("run-1", ws2)

    await manager.broadcast("run-1", {"event": "test"})

    assert ws1 not in manager._connections["run-1"]
    assert ws2 in manager._connections["run-1"]


async def test_emit_wraps_event(manager):
    ws = AsyncMock()
    await manager.connect("run-1", ws)

    await manager.emit("run-1", "case_start", {"case_id": "c1"})

    msg = json.loads(ws.send_text.call_args[0][0])
    assert msg["event"] == "case_start"
    assert msg["case_id"] == "c1"