from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock

import pytest

from yuxi.channel.interfaces.websocket.manager import WSManager


class TestWSManager:
    @pytest.fixture
    def manager(self):
        return WSManager()

    @pytest.mark.asyncio
    async def test_start(self, manager):
        await manager.start()
        assert manager.is_running is True

    @pytest.mark.asyncio
    async def test_stop(self, manager):
        await manager.start()
        await manager.stop()
        assert manager.is_running is False

    @pytest.mark.asyncio
    async def test_stop_all(self, manager):
        await manager.start()
        ws = AsyncMock()
        manager._connections["conn1"] = ws
        await manager.stop_all()
        ws.close.assert_awaited_once()
        assert len(manager._connections) == 0

    @pytest.mark.asyncio
    async def test_register(self, manager):
        await manager.start()
        ws = AsyncMock()
        await manager.register("conn1", ws)
        assert "conn1" in manager._connections

    @pytest.mark.asyncio
    async def test_unregister(self, manager):
        await manager.start()
        ws = AsyncMock()
        await manager.register("conn1", ws)
        await manager.unregister("conn1")
        assert "conn1" not in manager._connections

    @pytest.mark.asyncio
    async def test_send_to(self, manager):
        await manager.start()
        ws = AsyncMock()
        await manager.register("conn1", ws)
        await manager.send_to("conn1", {"type": "message"})
        ws.send.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_send_to_nonexistent(self, manager):
        await manager.start()
        await manager.send_to("nonexistent", {"type": "message"})

    @pytest.mark.asyncio
    async def test_broadcast(self, manager):
        await manager.start()
        ws1 = AsyncMock()
        ws2 = AsyncMock()
        await manager.register("conn1", ws1)
        await manager.register("conn2", ws2)
        await manager.broadcast({"type": "message"})
        ws1.send.assert_awaited_once()
        ws2.send.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_get_connection(self, manager):
        await manager.start()
        ws = AsyncMock()
        await manager.register("conn1", ws)
        result = manager.get_connection("conn1")
        assert result is ws

    def test_get_connection_not_found(self, manager):
        result = manager.get_connection("nonexistent")
        assert result is None

    def test_connection_count(self, manager):
        assert manager.connection_count == 0