"""Integration tests for WebSocket endpoints."""

from fastapi.testclient import TestClient
from unittest.mock import MagicMock, patch, AsyncMock

from preloop.models.models.user import User


class TestUnifiedWebSocket:
    """Test the unified WebSocket endpoint /api/v1/ws/unified."""

    def _setup_websocket_mocks(
        self, mock_manager, mock_get_user, mock_session_manager, test_user=None
    ):
        """Helper to setup async mocks for WebSocket tests."""
        # Mock user authentication (async function)
        async_mock_get_user = AsyncMock(return_value=test_user)
        mock_get_user.side_effect = async_mock_get_user

        # Mock session creation (async function)
        mock_session = MagicMock()
        mock_session.id = "test-session-id"
        mock_session.connection_id = "test-connection-id"
        mock_session.is_authenticated = test_user is not None
        async_mock_create_session = AsyncMock(return_value=mock_session)
        mock_session_manager.create_session = async_mock_create_session
        mock_session_manager.update_activity = MagicMock()

        # Mock manager connection (async function)
        async_mock_connect = AsyncMock(return_value="test-manager-connection-id")
        mock_manager.connect_with_account = async_mock_connect
        mock_manager.disconnect = MagicMock()
        mock_manager.active_connections = {}

        # Mock session cleanup (async function)
        async_mock_end_session = AsyncMock()
        mock_session_manager.end_session = async_mock_end_session

        return mock_session

    @patch("preloop.api.endpoints.websockets.handle_activity")
    @patch("preloop.api.endpoints.websockets.session_manager")
    @patch("preloop.api.endpoints.websockets.get_user_from_token_if_valid")
    @patch("preloop.api.endpoints.websockets.manager")
    def test_unified_websocket_handshake(
        self,
        mock_manager,
        mock_get_user,
        mock_session_manager,
        mock_handle_activity,
        client: TestClient,
        test_user: User,
    ):
        """Test WebSocket handshake with authenticated user."""
        mock_session = self._setup_websocket_mocks(
            mock_manager, mock_get_user, mock_session_manager, test_user
        )

        # Connect to WebSocket
        with client.websocket_connect(
            "/api/v1/ws/unified?token=test-token&fingerprint=test-fingerprint"
        ) as websocket:
            # Should receive handshake confirmation
            data = websocket.receive_json()
            assert data["type"] == "handshake"
            assert data["session_id"] == "test-session-id"
            assert data["authenticated"] is True
            assert data["message"] == "Connected to unified WebSocket"

    @patch("preloop.api.endpoints.websockets.handle_activity")
    @patch("preloop.api.endpoints.websockets.session_manager")
    @patch("preloop.api.endpoints.websockets.get_user_from_token_if_valid")
    @patch("preloop.api.endpoints.websockets.manager")
    def test_unified_websocket_ping_pong(
        self,
        mock_manager,
        mock_get_user,
        mock_session_manager,
        mock_handle_activity,
        client: TestClient,
        test_user: User,
    ):
        """Test WebSocket ping/pong heartbeat lifecycle."""
        mock_session = self._setup_websocket_mocks(
            mock_manager, mock_get_user, mock_session_manager, test_user
        )

        with client.websocket_connect(
            "/api/v1/ws/unified?token=test-token"
        ) as websocket:
            # Receive handshake
            handshake = websocket.receive_json()
            assert handshake["type"] == "handshake"

            # Send activity message (pong doesn't update activity, but other messages do)
            websocket.send_json({"type": "activity", "event_type": "test"})

            # Give it a moment to process
            import time

            time.sleep(0.05)

            # Verify session activity was updated
            mock_session_manager.update_activity.assert_called()

    @patch("preloop.api.endpoints.websockets.handle_activity")
    @patch("preloop.api.endpoints.websockets.session_manager")
    @patch("preloop.api.endpoints.websockets.manager")
    def test_unified_websocket_invalid_token(
        self,
        mock_manager,
        mock_session_manager,
        mock_handle_activity,
        client: TestClient,
    ):
        """Test WebSocket connection with invalid token results in anonymous session.

        The /ws/unified endpoint allows anonymous connections, so an invalid token
        should result in an anonymous session (not an error).
        """
        # Setup mocks for anonymous user
        mock_session = MagicMock()
        mock_session.id = "anon-session-id"
        mock_session.connection_id = "anon-connection-id"
        mock_session.is_authenticated = False
        async_mock_create_session = AsyncMock(return_value=mock_session)
        mock_session_manager.create_session = async_mock_create_session
        mock_session_manager.update_activity = MagicMock()
        async_mock_end_session = AsyncMock()
        mock_session_manager.end_session = async_mock_end_session

        mock_manager.active_connections = {}
        mock_manager.disconnect = MagicMock()

        # Connect with invalid token - should get anonymous session
        with client.websocket_connect(
            "/api/v1/ws/unified?token=invalid-token"
        ) as websocket:
            # Should receive handshake (anonymous allowed)
            data = websocket.receive_json()
            assert data["type"] == "handshake"
            assert data["authenticated"] is False  # Invalid token = anonymous

    @patch("preloop.api.endpoints.websockets.handle_activity")
    @patch("preloop.api.endpoints.websockets.session_manager")
    @patch("preloop.api.endpoints.websockets.get_user_from_token_if_valid")
    @patch("preloop.api.endpoints.websockets.manager")
    def test_unified_websocket_anonymous_connection(
        self,
        mock_manager,
        mock_get_user,
        mock_session_manager,
        mock_handle_activity,
        client: TestClient,
    ):
        """Test WebSocket connection without authentication (anonymous)."""
        # Setup mocks for anonymous user (test_user=None)
        # Mock user authentication (async function)
        async_mock_get_user = AsyncMock(return_value=None)
        mock_get_user.side_effect = async_mock_get_user

        # Mock session creation for anonymous user
        mock_session = MagicMock()
        mock_session.id = "anon-session-id"
        mock_session.connection_id = "anon-connection-id"
        mock_session.is_authenticated = False
        async_mock_create_session = AsyncMock(return_value=mock_session)
        mock_session_manager.create_session = async_mock_create_session
        mock_session_manager.update_activity = MagicMock()

        # Mock manager for anonymous connection
        mock_manager.active_connections = {}
        mock_manager.disconnect = MagicMock()

        # Mock session cleanup
        async_mock_end_session = AsyncMock()
        mock_session_manager.end_session = async_mock_end_session

        with client.websocket_connect(
            "/api/v1/ws/unified?fingerprint=anon-fingerprint"
        ) as websocket:
            # Should receive handshake for anonymous user
            data = websocket.receive_json()
            assert data["type"] == "handshake"
            assert data["session_id"] == "anon-session-id"
            assert data["authenticated"] is False

    @patch("preloop.api.endpoints.websockets.handle_activity")
    @patch("preloop.api.endpoints.websockets.session_manager")
    @patch("preloop.api.endpoints.websockets.get_user_from_token_if_valid")
    @patch("preloop.api.endpoints.websockets.manager")
    def test_unified_websocket_activity_tracking(
        self,
        mock_manager,
        mock_get_user,
        mock_session_manager,
        mock_handle_activity,
        client: TestClient,
        test_user: User,
    ):
        """Test activity tracking through WebSocket messages."""
        mock_session = self._setup_websocket_mocks(
            mock_manager, mock_get_user, mock_session_manager, test_user
        )

        # Mock handle_activity as async
        async_mock_handle_activity = AsyncMock()
        mock_handle_activity.side_effect = async_mock_handle_activity

        with client.websocket_connect(
            "/api/v1/ws/unified?token=test-token"
        ) as websocket:
            # Receive handshake
            handshake = websocket.receive_json()
            assert handshake["type"] == "handshake"

            # Send activity event
            activity_data = {
                "type": "activity",
                "event_type": "page_view",
                "path": "/dashboard",
            }
            websocket.send_json(activity_data)

            # Give it time to process
            import time

            time.sleep(0.1)

            # Verify activity handler was called
            mock_handle_activity.assert_called_once()

    @patch("preloop.api.endpoints.websockets.handle_activity")
    @patch("preloop.api.endpoints.websockets.session_manager")
    @patch("preloop.api.endpoints.websockets.manager")
    def test_unified_websocket_cleanup_on_disconnect(
        self,
        mock_manager,
        mock_session_manager,
        mock_handle_activity,
        client: TestClient,
    ):
        """Test proper cleanup when WebSocket disconnects (anonymous user)."""
        # Setup mocks for anonymous user
        mock_session = MagicMock()
        mock_session.id = "test-session-id"
        mock_session.connection_id = "test-connection-id"
        mock_session.is_authenticated = False
        async_mock_create_session = AsyncMock(return_value=mock_session)
        mock_session_manager.create_session = async_mock_create_session
        mock_session_manager.update_activity = MagicMock()
        async_mock_end_session = AsyncMock()
        mock_session_manager.end_session = async_mock_end_session

        mock_manager.active_connections = {}
        mock_manager.disconnect = MagicMock()

        with client.websocket_connect(
            "/api/v1/ws/unified?fingerprint=test-fingerprint"
        ) as websocket:
            # Receive handshake
            handshake = websocket.receive_json()
            assert handshake["type"] == "handshake"

            # Close connection
            websocket.close()

        # Verify cleanup was called
        mock_manager.disconnect.assert_called_once()
        mock_session_manager.end_session.assert_called_once()

    @patch("preloop.api.endpoints.websockets.handle_activity")
    @patch("preloop.api.endpoints.websockets.session_manager")
    @patch("preloop.api.endpoints.websockets.manager")
    def test_unified_websocket_cleanup_on_exception(
        self,
        mock_manager,
        mock_session_manager,
        mock_handle_activity,
        client: TestClient,
    ):
        """Test cleanup runs even if manager raises exception (regression test)."""
        # Mock session creation (async function)
        mock_session = MagicMock()
        mock_session.id = "test-session-id"
        mock_session.connection_id = "test-connection-id"
        mock_session.is_authenticated = False
        async_mock_create_session = AsyncMock(return_value=mock_session)
        mock_session_manager.create_session = async_mock_create_session

        # Mock manager to raise exception when setting active_connections
        mock_manager.active_connections = {}
        mock_manager.disconnect = MagicMock()

        # Mock session cleanup (async function)
        async_mock_end_session = AsyncMock()
        mock_session_manager.end_session = async_mock_end_session

        # Anonymous connection should work (no exception expected)
        with client.websocket_connect(
            "/api/v1/ws/unified?fingerprint=test-fingerprint"
        ) as websocket:
            # Should receive handshake
            handshake = websocket.receive_json()
            assert handshake["type"] == "handshake"
            assert handshake["authenticated"] is False

            # Close normally
            websocket.close()

        # Verify session cleanup was called
        mock_session_manager.end_session.assert_called_once()
