"""Tests for local WebSocket server."""

import pytest
import asyncio
import json
import tempfile
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock

from skop_runner.local_server import LocalWebSocketServer
from skop_runner.message_handler import MessageHandler


@pytest.fixture
def mock_message_handler():
    """Create a mock message handler."""
    handler = MagicMock(spec=MessageHandler)
    handler.handle_message = AsyncMock(return_value={
        "jsonrpc": "2.0",
        "id": "1",
        "result": {"root": str(Path(tempfile.gettempdir()) / "test")},
    })
    return handler


@pytest.fixture
def server(mock_message_handler):
    """Create a LocalWebSocketServer instance."""
    return LocalWebSocketServer(mock_message_handler)


def _ws_url(server):
    """Build the local WS URL with token."""
    return f"ws://127.0.0.1:{server.port}?token={server.token}"


class TestServerLifecycle:
    """Test server start/stop lifecycle."""

    @pytest.mark.asyncio
    async def test_start_returns_port(self, server):
        """start() should return an OS-assigned port > 0."""
        port = await server.start()
        try:
            assert port > 0
            assert server.port == port
        finally:
            await server.stop()

    @pytest.mark.asyncio
    async def test_port_is_none_before_start(self, server):
        """port property is None before calling start()."""
        assert server.port is None

    @pytest.mark.asyncio
    async def test_stop_clears_state(self, server):
        """stop() resets port and server references."""
        await server.start()
        await server.stop()

        assert server.port is None
        assert server._server is None
        assert len(server._clients) == 0

    @pytest.mark.asyncio
    async def test_stop_when_not_started(self, server):
        """stop() is safe to call before start()."""
        await server.stop()  # Should not raise

    @pytest.mark.asyncio
    async def test_token_generated_on_init(self, server):
        """A connection token should be generated on construction."""
        assert server.token
        assert len(server.token) > 16


class TestTokenAuth:
    """Test CSWSH protection via connection token."""

    @pytest.mark.asyncio
    async def test_connection_without_token_rejected(self, server):
        """Connections without a token should be closed immediately."""
        port = await server.start()

        try:
            import websockets

            ws = await websockets.connect(f"ws://127.0.0.1:{port}")
            # Server should close the connection
            await asyncio.wait_for(ws.wait_closed(), timeout=2.0)
            assert ws.close_code == 4003
            assert len(server._clients) == 0
        finally:
            await server.stop()

    @pytest.mark.asyncio
    async def test_connection_with_wrong_token_rejected(self, server):
        """Connections with an invalid token should be closed immediately."""
        port = await server.start()

        try:
            import websockets

            ws = await websockets.connect(
                f"ws://127.0.0.1:{port}?token=wrong-token"
            )
            await asyncio.wait_for(ws.wait_closed(), timeout=2.0)
            assert ws.close_code == 4003
            assert len(server._clients) == 0
        finally:
            await server.stop()

    @pytest.mark.asyncio
    async def test_connection_with_valid_token_accepted(self, server):
        """Connections with the correct token should succeed."""
        port = await server.start()

        try:
            import websockets

            async with websockets.connect(_ws_url(server)) as ws:
                await asyncio.sleep(0.05)
                assert len(server._clients) == 1
        finally:
            await server.stop()


class TestMessageHandling:
    """Test message routing through the local server."""

    @pytest.mark.asyncio
    async def test_handle_message_routes_to_handler(self, server, mock_message_handler):
        """Messages should be routed through MessageHandler."""
        port = await server.start()

        try:
            import websockets

            async with websockets.connect(_ws_url(server)) as ws:
                request = {
                    "jsonrpc": "2.0",
                    "id": "1",
                    "method": "workspace.getRoot",
                    "params": {},
                }
                await ws.send(json.dumps(request))
                response = json.loads(await ws.recv())

                assert response["id"] == "1"
                assert "result" in response
                mock_message_handler.handle_message.assert_called_once_with(request)
        finally:
            await server.stop()

    @pytest.mark.asyncio
    async def test_invalid_json_ignored(self, server, mock_message_handler):
        """Invalid JSON should be silently ignored."""
        port = await server.start()

        try:
            import websockets

            async with websockets.connect(_ws_url(server)) as ws:
                await ws.send("not valid json")
                # Send a valid message to confirm the connection is still alive
                request = {
                    "jsonrpc": "2.0",
                    "id": "2",
                    "method": "test",
                    "params": {},
                }
                await ws.send(json.dumps(request))
                response = json.loads(await ws.recv())

                assert response["id"] == "1"  # From mock handler
        finally:
            await server.stop()


class TestBroadcastNotification:
    """Test notification broadcasting."""

    @pytest.mark.asyncio
    async def test_broadcast_to_connected_clients(self, server):
        """Notifications should reach all connected local clients."""
        port = await server.start()

        try:
            import websockets

            async with websockets.connect(_ws_url(server)) as ws:
                # Wait for client to be registered
                await asyncio.sleep(0.05)

                await server._broadcast_notification(
                    "kernel.output", {"cell_id": "c1", "data": "hello"}
                )

                raw = await asyncio.wait_for(ws.recv(), timeout=2.0)
                notification = json.loads(raw)

                assert notification["jsonrpc"] == "2.0"
                assert notification["method"] == "kernel.output"
                assert notification["params"]["cell_id"] == "c1"
                assert "id" not in notification  # Notifications have no id
        finally:
            await server.stop()

    @pytest.mark.asyncio
    async def test_broadcast_no_clients_is_noop(self, server):
        """Broadcasting with no clients should not raise."""
        await server.start()
        try:
            await server._broadcast_notification("kernel.output", {"cell_id": "c1"})
        finally:
            await server.stop()

    @pytest.mark.asyncio
    async def test_broadcast_all_messages_sent(self, server):
        """All broadcast messages should be sent (rate limiting is in multiplexer)."""
        port = await server.start()

        try:
            import websockets

            async with websockets.connect(_ws_url(server)) as ws:
                await asyncio.sleep(0.05)

                count = 50
                for i in range(count):
                    await server._broadcast_notification(
                        "kernel.output", {"i": i}
                    )

                received = []
                try:
                    while True:
                        raw = await asyncio.wait_for(ws.recv(), timeout=0.5)
                        received.append(json.loads(raw))
                except asyncio.TimeoutError:
                    pass

                assert len(received) == count
        finally:
            await server.stop()


class TestClientDisconnect:
    """Test client disconnect handling."""

    @pytest.mark.asyncio
    async def test_client_disconnect_cleanup(self, server):
        """Disconnected clients should be removed from the set."""
        port = await server.start()

        try:
            import websockets

            ws = await websockets.connect(_ws_url(server))
            await asyncio.sleep(0.05)
            assert len(server._clients) == 1

            await ws.close()
            await asyncio.sleep(0.1)
            assert len(server._clients) == 0
        finally:
            await server.stop()

    @pytest.mark.asyncio
    async def test_multiple_clients(self, server):
        """Multiple clients can connect simultaneously."""
        port = await server.start()

        try:
            import websockets

            ws1 = await websockets.connect(_ws_url(server))
            ws2 = await websockets.connect(_ws_url(server))
            await asyncio.sleep(0.05)

            assert len(server._clients) == 2

            await ws1.close()
            await asyncio.sleep(0.1)
            assert len(server._clients) == 1

            await ws2.close()
            await asyncio.sleep(0.1)
            assert len(server._clients) == 0
        finally:
            await server.stop()

    @pytest.mark.asyncio
    async def test_in_flight_tasks_drained_on_disconnect(self, server, mock_message_handler):
        """In-flight handlers should be given time to complete on disconnect."""
        handler_started = asyncio.Event()
        handler_finished = asyncio.Event()

        async def slow_handler(message):
            handler_started.set()
            await asyncio.sleep(0.5)
            handler_finished.set()
            return {"jsonrpc": "2.0", "id": message.get("id"), "result": "ok"}

        mock_message_handler.handle_message = slow_handler
        port = await server.start()

        try:
            import websockets

            ws = await websockets.connect(_ws_url(server))
            await asyncio.sleep(0.05)

            # Send a message that will take 0.5s to handle
            await ws.send(json.dumps({
                "jsonrpc": "2.0", "id": "1", "method": "slow", "params": {}
            }))

            # Wait for handler to start, then disconnect
            await asyncio.wait_for(handler_started.wait(), timeout=2.0)
            await ws.close()

            # The handler should still complete (within the 2s drain timeout)
            await asyncio.wait_for(handler_finished.wait(), timeout=3.0)
        finally:
            await server.stop()
