"""
Tests for WebSocket manager.

Tests connection management, channel subscriptions, and message routing.
"""

from datetime import UTC, datetime, timedelta
from typing import Any
from unittest.mock import AsyncMock

import pytest

from dazzle_back.runtime.websocket_manager import (
    Connection,
    MessageType,
    RealtimeMessage,
    create_websocket_manager,
)

# =============================================================================
# Fixtures
# =============================================================================


@pytest.fixture
def ws_manager() -> Any:
    """Create a WebSocket manager for testing."""
    return create_websocket_manager(
        max_subscriptions=10,
        heartbeat_timeout=30,
    )


@pytest.fixture
def mock_websocket() -> Any:
    """Create a mock WebSocket."""
    ws = AsyncMock()
    ws.accept = AsyncMock()
    ws.send_json = AsyncMock()
    ws.receive_json = AsyncMock()
    ws.close = AsyncMock()
    return ws


# =============================================================================
# Connection Tests
# =============================================================================


class TestConnection:
    """Tests for Connection class."""

    def test_create_connection(self, mock_websocket: Any) -> None:
        """Test creating a connection."""
        conn = Connection(
            id="conn_123",
            websocket=mock_websocket,
            user_id="user_456",
            user_name="Alice",
        )

        assert conn.id == "conn_123"
        assert conn.user_id == "user_456"
        assert conn.user_name == "Alice"
        assert len(conn.subscriptions) == 0
        assert isinstance(conn.connected_at, datetime)
        assert isinstance(conn.last_heartbeat, datetime)

    def test_update_heartbeat(self, mock_websocket: Any) -> None:
        """Test updating heartbeat."""
        conn = Connection(id="conn_123", websocket=mock_websocket)
        old_heartbeat = conn.last_heartbeat

        # Wait a tiny bit then update
        conn.update_heartbeat()

        assert conn.last_heartbeat >= old_heartbeat


# =============================================================================
# RealtimeMessage Tests
# =============================================================================


class TestRealtimeMessage:
    """Tests for RealtimeMessage class."""

    def test_create_message(self) -> None:
        """Test creating a message."""
        msg = RealtimeMessage(
            type=MessageType.ENTITY_CREATED,
            channel="entity:Task",
            payload={"id": "123", "title": "Test"},
        )

        assert msg.type == MessageType.ENTITY_CREATED
        assert msg.channel == "entity:Task"
        assert msg.payload is not None and msg.payload["title"] == "Test"
        assert msg.timestamp > 0

    def test_to_dict(self) -> None:
        """Test converting message to dict."""
        msg = RealtimeMessage(
            type=MessageType.SUBSCRIBE,
            channel="entity:Task",
            request_id="req_123",
        )

        d = msg.to_dict()

        assert d["type"] == "subscribe"
        assert d["channel"] == "entity:Task"
        assert d["requestId"] == "req_123"
        assert "timestamp" in d

    def test_from_dict(self) -> None:
        """Test creating message from dict."""
        d = {
            "type": "entity:updated",
            "channel": "entity:Task",
            "payload": {"id": "123"},
            "requestId": "req_456",
            "timestamp": 1234567890000,
        }

        msg = RealtimeMessage.from_dict(d)

        assert msg.type == "entity:updated"
        assert msg.channel == "entity:Task"
        assert msg.payload is not None and msg.payload["id"] == "123"
        assert msg.request_id == "req_456"
        assert msg.timestamp == 1234567890000


# =============================================================================
# WebSocketManager Tests
# =============================================================================


class TestWebSocketManager:
    """Tests for WebSocketManager class."""

    @pytest.mark.asyncio
    async def test_connect(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test accepting a connection."""
        connection_id = await ws_manager.connect(
            mock_websocket,
            user_id="user_123",
            user_name="Alice",
        )

        assert connection_id is not None
        assert ws_manager.connection_count == 1

        mock_websocket.accept.assert_called_once()
        mock_websocket.send_json.assert_called_once()

        # Verify connected message
        call_args = mock_websocket.send_json.call_args[0][0]
        assert call_args["type"] == "connected"
        assert call_args["payload"]["connectionId"] == connection_id
        assert call_args["payload"]["userId"] == "user_123"

    @pytest.mark.asyncio
    async def test_disconnect(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test disconnecting."""
        connection_id = await ws_manager.connect(mock_websocket, user_id="user_123")
        assert ws_manager.connection_count == 1

        await ws_manager.disconnect(connection_id)

        assert ws_manager.connection_count == 0
        assert ws_manager.get_connection(connection_id) is None

    @pytest.mark.asyncio
    async def test_get_connection(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test getting a connection."""
        connection_id = await ws_manager.connect(mock_websocket, user_id="user_123")

        conn = ws_manager.get_connection(connection_id)

        assert conn is not None
        assert conn.id == connection_id
        assert conn.user_id == "user_123"

    @pytest.mark.asyncio
    async def test_get_user_connections(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test getting all connections for a user."""
        # Connect same user twice
        conn1 = await ws_manager.connect(mock_websocket, user_id="user_123")
        mock_websocket2 = AsyncMock()
        mock_websocket2.accept = AsyncMock()
        mock_websocket2.send_json = AsyncMock()
        conn2 = await ws_manager.connect(mock_websocket2, user_id="user_123")

        connections = ws_manager.get_user_connections("user_123")

        assert len(connections) == 2
        assert {c.id for c in connections} == {conn1, conn2}


# =============================================================================
# Channel Subscription Tests
# =============================================================================


class TestChannelSubscriptions:
    """Tests for channel subscriptions."""

    @pytest.mark.asyncio
    async def test_subscribe(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test subscribing to a channel."""
        connection_id = await ws_manager.connect(mock_websocket)

        result = await ws_manager.subscribe(connection_id, "entity:Task")

        assert result is True
        assert ws_manager.is_subscribed(connection_id, "entity:Task")
        assert ws_manager.channel_count == 1

        # Verify subscribed message
        assert mock_websocket.send_json.call_count == 2  # connected + subscribed
        last_call = mock_websocket.send_json.call_args[0][0]
        assert last_call["type"] == "subscribed"
        assert last_call["channel"] == "entity:Task"

    @pytest.mark.asyncio
    async def test_unsubscribe(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test unsubscribing from a channel."""
        connection_id = await ws_manager.connect(mock_websocket)
        await ws_manager.subscribe(connection_id, "entity:Task")

        result = await ws_manager.unsubscribe(connection_id, "entity:Task")

        assert result is True
        assert not ws_manager.is_subscribed(connection_id, "entity:Task")
        assert ws_manager.channel_count == 0

    @pytest.mark.asyncio
    async def test_subscription_limit(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test subscription limit is enforced."""
        connection_id = await ws_manager.connect(mock_websocket)

        # Subscribe to max channels
        for i in range(10):
            await ws_manager.subscribe(connection_id, f"channel:{i}")

        # Try to subscribe to one more
        result = await ws_manager.subscribe(connection_id, "channel:extra")

        assert result is False
        assert not ws_manager.is_subscribed(connection_id, "channel:extra")

    @pytest.mark.asyncio
    async def test_get_channel_subscribers(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test getting channel subscribers."""
        conn1 = await ws_manager.connect(mock_websocket)
        mock_websocket2 = AsyncMock()
        mock_websocket2.accept = AsyncMock()
        mock_websocket2.send_json = AsyncMock()
        conn2 = await ws_manager.connect(mock_websocket2)

        await ws_manager.subscribe(conn1, "entity:Task")
        await ws_manager.subscribe(conn2, "entity:Task")

        subscribers = ws_manager.get_channel_subscribers("entity:Task")

        assert len(subscribers) == 2
        assert set(subscribers) == {conn1, conn2}

    @pytest.mark.asyncio
    async def test_disconnect_removes_subscriptions(
        self, ws_manager: Any, mock_websocket: Any
    ) -> None:
        """Test that disconnect removes all subscriptions."""
        connection_id = await ws_manager.connect(mock_websocket)
        await ws_manager.subscribe(connection_id, "entity:Task")
        await ws_manager.subscribe(connection_id, "entity:User")

        await ws_manager.disconnect(connection_id)

        assert ws_manager.channel_count == 0


# =============================================================================
# Message Handling Tests
# =============================================================================


class TestMessageHandling:
    """Tests for message handling."""

    @pytest.mark.asyncio
    async def test_handle_ping(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test handling ping message."""
        connection_id = await ws_manager.connect(mock_websocket)
        mock_websocket.send_json.reset_mock()

        await ws_manager.handle_message(connection_id, {"type": "ping"})

        mock_websocket.send_json.assert_called_once()
        call_args = mock_websocket.send_json.call_args[0][0]
        assert call_args["type"] == "pong"

    @pytest.mark.asyncio
    async def test_handle_subscribe(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test handling subscribe message."""
        connection_id = await ws_manager.connect(mock_websocket)

        await ws_manager.handle_message(
            connection_id,
            {
                "type": "subscribe",
                "channel": "entity:Task",
                "requestId": "req_123",
            },
        )

        assert ws_manager.is_subscribed(connection_id, "entity:Task")

    @pytest.mark.asyncio
    async def test_handle_unsubscribe(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test handling unsubscribe message."""
        connection_id = await ws_manager.connect(mock_websocket)
        await ws_manager.subscribe(connection_id, "entity:Task")

        await ws_manager.handle_message(
            connection_id,
            {
                "type": "unsubscribe",
                "channel": "entity:Task",
            },
        )

        assert not ws_manager.is_subscribed(connection_id, "entity:Task")

    @pytest.mark.asyncio
    async def test_handle_unknown_message(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test handling unknown message type."""
        connection_id = await ws_manager.connect(mock_websocket)
        mock_websocket.send_json.reset_mock()

        await ws_manager.handle_message(
            connection_id,
            {
                "type": "unknown_type",
                "requestId": "req_123",
            },
        )

        # Should send error
        mock_websocket.send_json.assert_called_once()
        call_args = mock_websocket.send_json.call_args[0][0]
        assert call_args["type"] == "error"
        assert "UNKNOWN_MESSAGE_TYPE" in call_args["payload"]["code"]

    @pytest.mark.asyncio
    async def test_custom_handler(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test registering and using custom handler."""
        handler_called = []

        async def custom_handler(conn_id: Any, message: Any) -> None:
            handler_called.append((conn_id, message.type))

        ws_manager.register_handler("custom:event", custom_handler)
        connection_id = await ws_manager.connect(mock_websocket)

        await ws_manager.handle_message(
            connection_id,
            {
                "type": "custom:event",
                "payload": {"data": "test"},
            },
        )

        assert len(handler_called) == 1
        assert handler_called[0][0] == connection_id
        assert handler_called[0][1] == "custom:event"


# =============================================================================
# Broadcasting Tests
# =============================================================================


class TestBroadcasting:
    """Tests for message broadcasting."""

    @pytest.mark.asyncio
    async def test_broadcast_to_channel(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test broadcasting to channel subscribers."""
        conn1 = await ws_manager.connect(mock_websocket)
        mock_websocket2 = AsyncMock()
        mock_websocket2.accept = AsyncMock()
        mock_websocket2.send_json = AsyncMock()
        conn2 = await ws_manager.connect(mock_websocket2)

        await ws_manager.subscribe(conn1, "entity:Task")
        await ws_manager.subscribe(conn2, "entity:Task")

        mock_websocket.send_json.reset_mock()
        mock_websocket2.send_json.reset_mock()

        message = RealtimeMessage(
            type=MessageType.ENTITY_CREATED,
            payload={"id": "123"},
        )
        sent_count = await ws_manager.broadcast("entity:Task", message)

        assert sent_count == 2
        assert mock_websocket.send_json.called
        assert mock_websocket2.send_json.called

    @pytest.mark.asyncio
    async def test_broadcast_exclude_sender(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test broadcasting with sender excluded."""
        conn1 = await ws_manager.connect(mock_websocket)
        mock_websocket2 = AsyncMock()
        mock_websocket2.accept = AsyncMock()
        mock_websocket2.send_json = AsyncMock()
        conn2 = await ws_manager.connect(mock_websocket2)

        await ws_manager.subscribe(conn1, "entity:Task")
        await ws_manager.subscribe(conn2, "entity:Task")

        mock_websocket.send_json.reset_mock()
        mock_websocket2.send_json.reset_mock()

        message = RealtimeMessage(type=MessageType.ENTITY_UPDATED)
        sent_count = await ws_manager.broadcast("entity:Task", message, exclude_connection=conn1)

        assert sent_count == 1
        assert not mock_websocket.send_json.called
        assert mock_websocket2.send_json.called

    @pytest.mark.asyncio
    async def test_broadcast_to_all(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test broadcasting to all connections."""
        await ws_manager.connect(mock_websocket)
        mock_websocket2 = AsyncMock()
        mock_websocket2.accept = AsyncMock()
        mock_websocket2.send_json = AsyncMock()
        await ws_manager.connect(mock_websocket2)

        mock_websocket.send_json.reset_mock()
        mock_websocket2.send_json.reset_mock()

        message = RealtimeMessage(type=MessageType.PING)
        sent_count = await ws_manager.broadcast_to_all(message)

        assert sent_count == 2

    @pytest.mark.asyncio
    async def test_send_to_user(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test sending to specific user."""
        await ws_manager.connect(mock_websocket, user_id="user_123")
        mock_websocket2 = AsyncMock()
        mock_websocket2.accept = AsyncMock()
        mock_websocket2.send_json = AsyncMock()
        await ws_manager.connect(mock_websocket2, user_id="user_456")

        mock_websocket.send_json.reset_mock()
        mock_websocket2.send_json.reset_mock()

        message = RealtimeMessage(type=MessageType.ENTITY_UPDATED)
        sent_count = await ws_manager.send_to_user("user_123", message)

        assert sent_count == 1
        assert mock_websocket.send_json.called
        assert not mock_websocket2.send_json.called


# =============================================================================
# Stale Connection Cleanup Tests
# =============================================================================


class TestStaleCleanup:
    """Tests for stale connection cleanup."""

    @pytest.mark.asyncio
    async def test_cleanup_stale_connections(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test cleaning up stale connections."""
        connection_id = await ws_manager.connect(mock_websocket)

        # Manually set heartbeat to past
        conn = ws_manager.get_connection(connection_id)
        conn.last_heartbeat = datetime.now(UTC) - timedelta(seconds=120)

        stale = await ws_manager.cleanup_stale_connections()

        assert len(stale) == 1
        assert stale[0] == connection_id
        assert ws_manager.connection_count == 0

    @pytest.mark.asyncio
    async def test_cleanup_keeps_active_connections(
        self, ws_manager: Any, mock_websocket: Any
    ) -> None:
        """Test that active connections are not cleaned up."""
        await ws_manager.connect(mock_websocket)

        # Connection is fresh, should not be cleaned up
        stale = await ws_manager.cleanup_stale_connections()

        assert len(stale) == 0
        assert ws_manager.connection_count == 1


# =============================================================================
# Statistics Tests
# =============================================================================


class TestStatistics:
    """Tests for manager statistics."""

    @pytest.mark.asyncio
    async def test_get_stats(self, ws_manager: Any, mock_websocket: Any) -> None:
        """Test getting statistics."""
        conn1 = await ws_manager.connect(mock_websocket, user_id="user_123")
        mock_websocket2 = AsyncMock()
        mock_websocket2.accept = AsyncMock()
        mock_websocket2.send_json = AsyncMock()
        conn2 = await ws_manager.connect(mock_websocket2, user_id="user_456")

        await ws_manager.subscribe(conn1, "entity:Task")
        await ws_manager.subscribe(conn1, "entity:User")
        await ws_manager.subscribe(conn2, "entity:Task")

        stats = ws_manager.get_stats()

        assert stats["connections"] == 2
        assert stats["channels"] == 2  # entity:Task and entity:User
        assert stats["users"] == 2
        assert stats["subscriptions"] == 3
