"""Tests for WebSocket client."""

import pytest
import asyncio
import json
from unittest.mock import AsyncMock, MagicMock, patch

from skop_runner.client import RunnerClient
from skop_runner.config import RunnerConfig
from skop_runner.message_handler import MessageHandler
from skop_runner.output_buffer import OutputBuffer


@pytest.fixture
def mock_config():
    """Create a mock config."""
    config = MagicMock(spec=RunnerConfig)
    config.ws_url = "ws://localhost:8080/runner"
    config.api_url = "http://localhost:8080"
    config.device_token = "test-token-123"
    return config


@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": {"success": True}
    })
    return handler


@pytest.fixture
def client(mock_config, mock_message_handler):
    """Create a RunnerClient instance."""
    return RunnerClient(mock_config, mock_message_handler, OutputBuffer())


class TestClientInit:
    """Test client initialization."""

    def test_init_defaults(self, client, mock_config, mock_message_handler):
        """Test client initializes with correct defaults."""
        assert client.config == mock_config
        assert client.message_handler == mock_message_handler
        assert client.ws is None
        assert client.running is False
        assert client._reconnect_delay == 1
        assert client._max_reconnect_delay == 15
        assert client.on_connect is None


class TestConnect:
    """Test connection handling."""

    @pytest.mark.asyncio
    async def test_connect_success(self, client, mock_config):
        """Test successful connection."""
        mock_ws = AsyncMock()

        with patch("skop_runner.client.websockets.connect", new_callable=AsyncMock) as mock_connect:
            mock_connect.return_value = mock_ws

            result = await client.connect()

            assert result is True
            assert client.running is True
            assert client.ws == mock_ws
            assert client._reconnect_delay == 1  # Reset on success

            # Verify connection URL
            expected_url = f"{mock_config.ws_url}?token={mock_config.device_token}"
            mock_connect.assert_called_once_with(
                expected_url,
                ping_interval=15,
                ping_timeout=20,
                max_size=50 * 1024 * 1024,
            )

    @pytest.mark.asyncio
    async def test_connect_failure(self, client):
        """Test connection failure."""
        with patch("skop_runner.client.websockets.connect", new_callable=AsyncMock) as mock_connect:
            mock_connect.side_effect = Exception("Connection refused")

            result = await client.connect()

            assert result is False
            assert client.running is False

    @pytest.mark.asyncio
    async def test_connect_resets_reconnect_delay(self, client):
        """Test that successful connection resets reconnect delay."""
        client._reconnect_delay = 30  # Simulate previous backoff

        with patch("skop_runner.client.websockets.connect", new_callable=AsyncMock) as mock_connect:
            mock_connect.return_value = AsyncMock()

            await client.connect()

            assert client._reconnect_delay == 1


class TestDisconnect:
    """Test disconnection handling."""

    @pytest.mark.asyncio
    async def test_disconnect(self, client):
        """Test disconnection."""
        mock_ws = AsyncMock()
        client.ws = mock_ws
        client.running = True

        await client.disconnect()

        assert client.running is False
        mock_ws.close.assert_called_once()

    @pytest.mark.asyncio
    async def test_disconnect_no_connection(self, client):
        """Test disconnection when not connected."""
        client.running = True

        await client.disconnect()

        assert client.running is False

    @pytest.mark.asyncio
    async def test_disconnect_close_error(self, client):
        """Test disconnection handles close errors gracefully."""
        mock_ws = AsyncMock()
        mock_ws.close.side_effect = Exception("Close error")
        client.ws = mock_ws
        client.running = True

        # Should not raise
        await client.disconnect()

        assert client.running is False


class TestReceiveLoop:
    """Test message receive loop."""

    @pytest.mark.asyncio
    async def test_receive_loop_not_connected(self, client):
        """Test receive loop when not connected raises error."""
        with pytest.raises(RuntimeError, match="Not connected"):
            await client.receive_loop()

    @pytest.mark.asyncio
    async def test_receive_loop_handles_messages(self, client, mock_message_handler):
        """Test receive loop processes messages correctly."""
        mock_ws = AsyncMock()
        messages = [
            json.dumps({"jsonrpc": "2.0", "method": "test.method", "id": "1", "params": {}}),
        ]
        mock_ws.__aiter__.return_value = iter(messages)
        client.ws = mock_ws
        client.running = True  # Normally set by connect()

        await client.receive_loop()

        mock_message_handler.handle_message.assert_called_once()
        mock_ws.send.assert_called_once()

    @pytest.mark.asyncio
    async def test_receive_loop_handles_invalid_json(self, client, mock_message_handler):
        """Test receive loop handles invalid JSON gracefully."""
        mock_ws = AsyncMock()
        messages = ["not valid json"]
        mock_ws.__aiter__.return_value = iter(messages)
        client.ws = mock_ws

        # Should not raise
        await client.receive_loop()

        # Handler should not be called for invalid JSON
        mock_message_handler.handle_message.assert_not_called()

    @pytest.mark.asyncio
    async def test_receive_loop_handles_handler_error(self, client, mock_message_handler):
        """Test receive loop handles handler errors gracefully."""
        mock_ws = AsyncMock()
        messages = [
            json.dumps({"jsonrpc": "2.0", "method": "test", "id": "1"}),
        ]
        mock_ws.__aiter__.return_value = iter(messages)
        mock_message_handler.handle_message.side_effect = Exception("Handler error")
        client.ws = mock_ws

        # Should not raise
        await client.receive_loop()


class TestHeartbeatLoop:
    """Test heartbeat loop."""

    @pytest.mark.asyncio
    async def test_heartbeat_sends_request(self, client, mock_config):
        """Test heartbeat sends POST request."""
        client.running = True

        # Create mock client and inject it directly
        mock_http_client = AsyncMock()
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_http_client.post.return_value = mock_response
        client._http_client = mock_http_client

        # Run heartbeat once then stop
        async def stop_after_heartbeat():
            await asyncio.sleep(0.1)
            client.running = False

        await asyncio.gather(
            client.heartbeat_loop(),
            stop_after_heartbeat(),
        )

        mock_http_client.post.assert_called_with(
            f"{mock_config.api_url}/api/runner/heartbeat",
            json={"token": mock_config.device_token},
        )

    @pytest.mark.asyncio
    async def test_heartbeat_handles_failure(self, client):
        """Test heartbeat handles request failure gracefully."""
        client.running = True

        # Create mock client and inject it directly
        mock_http_client = AsyncMock()
        mock_response = MagicMock()
        mock_response.status_code = 500
        mock_http_client.post.return_value = mock_response
        client._http_client = mock_http_client

        async def stop_after_heartbeat():
            await asyncio.sleep(0.1)
            client.running = False

        # Should not raise
        await asyncio.gather(
            client.heartbeat_loop(),
            stop_after_heartbeat(),
        )

    @pytest.mark.asyncio
    async def test_heartbeat_handles_exception(self, client):
        """Test heartbeat handles exceptions gracefully."""
        client.running = True

        # Create mock client that raises an exception
        mock_http_client = AsyncMock()
        mock_http_client.post.side_effect = Exception("Network error")
        client._http_client = mock_http_client

        async def stop_after_heartbeat():
            await asyncio.sleep(0.1)
            client.running = False

        # Should not raise
        await asyncio.gather(
            client.heartbeat_loop(),
            stop_after_heartbeat(),
        )

    @pytest.mark.asyncio
    async def test_heartbeat_includes_metadata(self, client, mock_config):
        """Test heartbeat includes metadata from supplier."""
        client.running = True
        client._metadata_supplier = lambda: {"local_ws_port": 12345}

        mock_http_client = AsyncMock()
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_http_client.post.return_value = mock_response
        client._http_client = mock_http_client

        async def stop_after_heartbeat():
            await asyncio.sleep(0.1)
            client.running = False

        await asyncio.gather(
            client.heartbeat_loop(),
            stop_after_heartbeat(),
        )

        mock_http_client.post.assert_called_with(
            f"{mock_config.api_url}/api/runner/heartbeat",
            json={
                "token": mock_config.device_token,
                "metadata": {"local_ws_port": 12345},
            },
        )

    @pytest.mark.asyncio
    async def test_heartbeat_no_metadata_when_supplier_not_set(self, client, mock_config):
        """Test heartbeat omits metadata when no supplier is set."""
        client.running = True
        assert client._metadata_supplier is None

        mock_http_client = AsyncMock()
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_http_client.post.return_value = mock_response
        client._http_client = mock_http_client

        async def stop_after_heartbeat():
            await asyncio.sleep(0.1)
            client.running = False

        await asyncio.gather(
            client.heartbeat_loop(),
            stop_after_heartbeat(),
        )

        mock_http_client.post.assert_called_with(
            f"{mock_config.api_url}/api/runner/heartbeat",
            json={"token": mock_config.device_token},
        )

    @pytest.mark.asyncio
    async def test_heartbeat_metadata_supplier_error(self, client, mock_config):
        """Test heartbeat gracefully handles metadata supplier errors."""
        client.running = True
        client._metadata_supplier = lambda: (_ for _ in ()).throw(RuntimeError("boom"))

        mock_http_client = AsyncMock()
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_http_client.post.return_value = mock_response
        client._http_client = mock_http_client

        async def stop_after_heartbeat():
            await asyncio.sleep(0.1)
            client.running = False

        await asyncio.gather(
            client.heartbeat_loop(),
            stop_after_heartbeat(),
        )

        # Should send without metadata when supplier raises
        mock_http_client.post.assert_called_with(
            f"{mock_config.api_url}/api/runner/heartbeat",
            json={"token": mock_config.device_token},
        )


class TestRunWithReconnect:
    """Test run with reconnect functionality."""

    @pytest.mark.asyncio
    async def test_run_with_reconnect_sets_running(self, client):
        """Test that run_with_reconnect sets running=True."""
        client.connect = AsyncMock(return_value=False)

        async def stop_client():
            await asyncio.sleep(0.1)
            client.running = False

        await asyncio.gather(
            client.run_with_reconnect(),
            stop_client(),
        )

        # Verify connect was called (loop entered)
        client.connect.assert_called()

    @pytest.mark.asyncio
    async def test_run_with_reconnect_exponential_backoff(self, client):
        """Test reconnection uses exponential backoff."""
        connect_count = 0

        async def mock_connect():
            nonlocal connect_count
            connect_count += 1
            if connect_count >= 3:
                client.running = False
            return False

        client.connect = mock_connect

        with patch("skop_runner.client.random.uniform", return_value=0):  # deterministic jitter
            with patch("skop_runner.client.asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
                await client.run_with_reconnect()

                # Check backoff delays: 1, 2 (then stops) — jitter is 0
                assert mock_sleep.call_count == 2
                mock_sleep.assert_any_call(1)
                mock_sleep.assert_any_call(2)

    @pytest.mark.asyncio
    async def test_run_with_reconnect_max_backoff(self, client):
        """Test reconnection respects max backoff."""
        client._reconnect_delay = 15  # Already at max

        async def mock_connect():
            client.running = False
            return False

        client.connect = mock_connect

        with patch("skop_runner.client.random.uniform", return_value=0):
            with patch("skop_runner.client.asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
                await client.run_with_reconnect()

                # Should not exceed max
                assert client._reconnect_delay == 15

    @pytest.mark.asyncio
    async def test_run_with_reconnect_success_runs_loops(self, client):
        """Test successful connection runs receive and heartbeat loops."""
        client.connect = AsyncMock(return_value=True)
        client.receive_loop = AsyncMock()
        client.heartbeat_loop = AsyncMock()

        async def stop_after_connect():
            await asyncio.sleep(0.1)
            client.running = False

        # Make loops exit quickly
        client.receive_loop.side_effect = stop_after_connect

        await client.run_with_reconnect()

        client.receive_loop.assert_called()
        client.heartbeat_loop.assert_called()

    @pytest.mark.asyncio
    async def test_run_with_reconnect_calls_on_connect_callback(self, client):
        """Test that on_connect callback is called on successful connection."""
        callback_called = []

        def on_connect():
            callback_called.append(True)

        client.on_connect = on_connect
        client.connect = AsyncMock(return_value=True)
        client.receive_loop = AsyncMock()
        client.heartbeat_loop = AsyncMock()

        async def stop_after_connect():
            await asyncio.sleep(0.1)
            client.running = False

        client.receive_loop.side_effect = stop_after_connect

        await client.run_with_reconnect()

        assert len(callback_called) == 1

    @pytest.mark.asyncio
    async def test_run_with_reconnect_no_callback_on_failed_connect(self, client):
        """Test that on_connect callback is not called on failed connection."""
        callback_called = []

        def on_connect():
            callback_called.append(True)

        client.on_connect = on_connect

        connect_count = 0

        async def mock_connect():
            nonlocal connect_count
            connect_count += 1
            if connect_count >= 2:
                client.running = False
            return False

        client.connect = mock_connect

        with patch("skop_runner.client.asyncio.sleep", new_callable=AsyncMock):
            await client.run_with_reconnect()

        assert len(callback_called) == 0


class TestOutputBufferIntegration:
    """Test output buffering when disconnected."""

    @pytest.mark.asyncio
    async def test_notification_buffered_when_disconnected(self, client):
        """Buffered methods are stored when client is not connected."""
        # Client starts with ws=None, running=False → disconnected
        assert client.ws is None
        await client._send_notification("kernel.output", {"cell_id": "c1"})
        assert len(client.output_buffer) == 1
        result = client.output_buffer.drain()
        assert result[0]["method"] == "kernel.output"
        assert result[0]["params"]["cell_id"] == "c1"

    @pytest.mark.asyncio
    async def test_notification_sent_when_connected(self, client):
        """Notifications are sent over WebSocket when connected."""
        mock_ws = AsyncMock()
        client.ws = mock_ws
        client.running = True
        await client._send_notification("kernel.output", {"cell_id": "c1"})
        # Should be sent over WS, not buffered
        mock_ws.send.assert_called_once()
        assert len(client.output_buffer) == 0

    @pytest.mark.asyncio
    async def test_non_buffered_method_dropped_when_disconnected(self, client):
        """Non-buffered methods are dropped silently when disconnected."""
        await client._send_notification("file.changed", {"path": "foo.py"})
        assert len(client.output_buffer) == 0

    @pytest.mark.asyncio
    async def test_notification_buffered_on_send_failure(self, client):
        """Buffered methods are stored when ws.send() raises."""
        mock_ws = AsyncMock()
        mock_ws.send.side_effect = Exception("Connection lost")
        client.ws = mock_ws
        client.running = True
        await client._send_notification("kernel.output", {"cell_id": "c1"})
        # Should be buffered after send failure
        assert len(client.output_buffer) == 1
        result = client.output_buffer.drain()
        assert result[0]["method"] == "kernel.output"
