"""Tests for kernel reconnection on runner restart (Phase 5.3)."""

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

from skop_runner.jupyter_backend.backend import JupyterBackend


class MockWebSocket:
    """Minimal mock for JupyterKernelWebSocket."""

    def __init__(self, **kwargs):
        self.kernel_id = kwargs.get("kernel_id")
        self.connected = False
        self.awaiting_input = False

    async def connect(self):
        self.connected = True

    async def close(self):
        self.connected = False


class TestReconnectSurvivingKernels:
    """Tests for _reconnect_surviving_kernels."""

    @pytest.fixture
    def backend(self):
        """Create a JupyterBackend with mocked internals."""
        b = JupyterBackend()
        # Mock the server manager
        b._server_manager = MagicMock()
        b._server_manager.base_url = "http://127.0.0.1:18888"
        b._server_manager.token = "test-token"
        b._server_manager.is_running = MagicMock(return_value=True)
        # Mock the API client
        b._api_client = AsyncMock()
        return b

    @pytest.mark.asyncio
    async def test_reconnect_surviving_kernel(self, backend):
        """Reconnect to a kernel that survived runner restart."""
        # Server has one kernel
        backend._api_client.list_kernels = AsyncMock(return_value=[
            {"id": "kernel-abc", "execution_state": "idle"},
        ])
        # That kernel belongs to a session
        backend._api_client.list_sessions = AsyncMock(return_value=[
            {"id": "session-xyz", "kernel": {"id": "kernel-abc"}},
        ])

        with patch(
            "skop_runner.jupyter_backend.backend.JupyterKernelWebSocket",
            side_effect=lambda **kw: MockWebSocket(**kw),
        ):
            count = await backend._reconnect_surviving_kernels()

        assert count == 1
        assert "kernel-abc" in backend._connections
        conn = backend._connections["kernel-abc"]
        assert conn.session_id == "session-xyz"
        assert conn.websocket.connected

    @pytest.mark.asyncio
    async def test_reconnect_skips_already_tracked(self, backend):
        """Already-tracked kernels should not be reconnected."""
        # Pre-populate a connection
        existing_conn = MagicMock()
        backend._connections["kernel-abc"] = existing_conn

        backend._api_client.list_kernels = AsyncMock(return_value=[
            {"id": "kernel-abc", "execution_state": "idle"},
        ])

        count = await backend._reconnect_surviving_kernels()

        assert count == 0
        # Original connection preserved
        assert backend._connections["kernel-abc"] is existing_conn

    @pytest.mark.asyncio
    async def test_reconnect_cleans_orphaned_kernel(self, backend):
        """Kernels with no session should be deleted."""
        backend._api_client.list_kernels = AsyncMock(return_value=[
            {"id": "kernel-orphan", "execution_state": "idle"},
        ])
        # No sessions at all
        backend._api_client.list_sessions = AsyncMock(return_value=[])

        count = await backend._reconnect_surviving_kernels()

        assert count == 0
        assert "kernel-orphan" not in backend._connections
        backend._api_client.delete_kernel.assert_called_once_with("kernel-orphan")

    @pytest.mark.asyncio
    async def test_reconnect_handles_ws_failure(self, backend):
        """If WebSocket connection fails, kernel should be cleaned up."""
        backend._api_client.list_kernels = AsyncMock(return_value=[
            {"id": "kernel-bad", "execution_state": "idle"},
        ])
        backend._api_client.list_sessions = AsyncMock(return_value=[
            {"id": "session-bad", "kernel": {"id": "kernel-bad"}},
        ])

        # WebSocket connect raises
        def bad_ws(**kw):
            ws = MockWebSocket(**kw)
            ws.connect = AsyncMock(side_effect=Exception("Connection failed"))
            return ws

        with patch(
            "skop_runner.jupyter_backend.backend.JupyterKernelWebSocket",
            side_effect=bad_ws,
        ):
            count = await backend._reconnect_surviving_kernels()

        assert count == 0
        assert "kernel-bad" not in backend._connections
        backend._api_client.delete_kernel.assert_called_once_with("kernel-bad")

    @pytest.mark.asyncio
    async def test_reconnect_multiple_kernels(self, backend):
        """Multiple surviving kernels should all be reconnected."""
        backend._api_client.list_kernels = AsyncMock(return_value=[
            {"id": "kernel-1", "execution_state": "idle"},
            {"id": "kernel-2", "execution_state": "busy"},
        ])
        backend._api_client.list_sessions = AsyncMock(return_value=[
            {"id": "session-1", "kernel": {"id": "kernel-1"}},
            {"id": "session-2", "kernel": {"id": "kernel-2"}},
        ])

        with patch(
            "skop_runner.jupyter_backend.backend.JupyterKernelWebSocket",
            side_effect=lambda **kw: MockWebSocket(**kw),
        ):
            count = await backend._reconnect_surviving_kernels()

        assert count == 2
        assert "kernel-1" in backend._connections
        assert "kernel-2" in backend._connections

    @pytest.mark.asyncio
    async def test_reconnect_no_api_client(self, backend):
        """Should return 0 when API client is not available."""
        backend._api_client = None
        count = await backend._reconnect_surviving_kernels()
        assert count == 0

    @pytest.mark.asyncio
    async def test_reconnect_handles_list_failure(self, backend):
        """Should return 0 if listing kernels fails."""
        backend._api_client.list_kernels = AsyncMock(
            side_effect=Exception("Server unreachable")
        )
        count = await backend._reconnect_surviving_kernels()
        assert count == 0
