"""Tests for Jupyter backend implementation."""

import pytest
from typing import Optional, Dict, Any, List, AsyncIterator

from skop_runner.kernel_backend import (
    KernelBackend,
    KernelInfo,
    KernelStatus,
    KernelNotFoundError,
    KernelStartError,
    register_backend,
    get_backend,
    list_backends,
    unregister_backend,
)


class TestJupyterBackendImports:
    """Test that jupyter_backend package imports correctly."""

    def test_package_imports(self):
        """All expected classes should be importable from package."""
        from skop_runner.jupyter_backend import (
            JupyterBackend,
            JupyterServerManager,
            JupyterAPIClient,
            JupyterKernelWebSocket,
        )
        assert JupyterBackend is not None
        assert JupyterServerManager is not None
        assert JupyterAPIClient is not None
        assert JupyterKernelWebSocket is not None

    def test_backend_module_imports(self):
        """Backend module should import without errors."""
        from skop_runner.jupyter_backend.backend import JupyterBackend
        assert JupyterBackend is not None

    def test_server_manager_imports(self):
        """Server manager should import without errors."""
        from skop_runner.jupyter_backend.server_manager import JupyterServerManager
        assert JupyterServerManager is not None

    def test_api_client_imports(self):
        """API client should import without errors."""
        from skop_runner.jupyter_backend.api_client import JupyterAPIClient
        assert JupyterAPIClient is not None

    def test_ws_client_imports(self):
        """WebSocket client should import without errors."""
        from skop_runner.jupyter_backend.ws_client import JupyterKernelWebSocket
        assert JupyterKernelWebSocket is not None


class TestJupyterBackendProtocolCompliance:
    """Verify JupyterBackend implements KernelBackend protocol."""

    def test_implements_protocol(self):
        """JupyterBackend should satisfy KernelBackend protocol."""
        from skop_runner.jupyter_backend import JupyterBackend
        backend = JupyterBackend()
        assert isinstance(backend, KernelBackend)

    def test_has_lifecycle_methods(self):
        """JupyterBackend should have all lifecycle methods."""
        from skop_runner.jupyter_backend import JupyterBackend

        methods = [
            "start_kernel",
            "stop_kernel",
            "restart_kernel",
            "interrupt_kernel",
            "shutdown_all",
        ]
        backend = JupyterBackend()
        for method in methods:
            assert hasattr(backend, method), f"Missing method: {method}"
            assert callable(getattr(backend, method)), f"{method} not callable"

    def test_has_execution_methods(self):
        """JupyterBackend should have all execution methods."""
        from skop_runner.jupyter_backend import JupyterBackend

        methods = ["execute_code", "send_input_reply"]
        backend = JupyterBackend()
        for method in methods:
            assert hasattr(backend, method), f"Missing method: {method}"
            assert callable(getattr(backend, method)), f"{method} not callable"

    def test_has_status_methods(self):
        """JupyterBackend should have all status methods."""
        from skop_runner.jupyter_backend import JupyterBackend

        methods = [
            "get_kernel_status",
            "get_kernel_info",
            "get_kernel_connection_info",
            "list_kernels",
            "list_available_kernels",
            "get_execution_count",
            "validate_kernel",
        ]
        backend = JupyterBackend()
        for method in methods:
            assert hasattr(backend, method), f"Missing method: {method}"
            assert callable(getattr(backend, method)), f"{method} not callable"

    def test_has_state_methods(self):
        """JupyterBackend should have all state methods."""
        from skop_runner.jupyter_backend import JupyterBackend

        methods = ["is_ready", "set_root_directory"]
        backend = JupyterBackend()
        for method in methods:
            assert hasattr(backend, method), f"Missing method: {method}"
            assert callable(getattr(backend, method)), f"{method} not callable"

    def test_has_lifecycle_hooks(self):
        """JupyterBackend should have all lifecycle hooks."""
        from skop_runner.jupyter_backend import JupyterBackend

        methods = ["initialize", "close"]
        backend = JupyterBackend()
        for method in methods:
            assert hasattr(backend, method), f"Missing method: {method}"
            assert callable(getattr(backend, method)), f"{method} not callable"


class TestJupyterBackendRegistration:
    """Test backend registration with factory."""

    def test_jupyter_backend_is_registered(self):
        """JupyterBackend should be registered on import."""
        # Importing the module should register the backend
        from skop_runner.jupyter_backend import JupyterBackend  # noqa: F401

        assert "jupyter" in list_backends()

    def test_get_jupyter_backend(self):
        """Should be able to get JupyterBackend via factory."""
        from skop_runner.jupyter_backend import JupyterBackend  # noqa: F401

        backend = get_backend("jupyter")
        assert isinstance(backend, JupyterBackend)

    def test_get_backend_with_callbacks(self):
        """Should be able to pass callbacks to backend constructor."""
        from skop_runner.jupyter_backend import JupyterBackend  # noqa: F401

        callback_called = []

        def on_died():
            callback_called.append("died")

        backend = get_backend("jupyter", on_backend_died=on_died)
        assert isinstance(backend, JupyterBackend)


class TestJupyterBackendUnitTests:
    """Unit tests for JupyterBackend (no Jupyter Server required)."""

    def test_not_ready_before_initialize(self):
        """Backend should not be ready before initialize() is called."""
        from skop_runner.jupyter_backend import JupyterBackend

        backend = JupyterBackend()
        assert backend.is_ready() is False

    def test_list_kernels_empty_initially(self):
        """list_kernels should return empty list initially."""
        from skop_runner.jupyter_backend import JupyterBackend

        backend = JupyterBackend()
        assert backend.list_kernels() == []

    def test_get_kernel_info_not_found(self):
        """get_kernel_info should return None for unknown kernel."""
        from skop_runner.jupyter_backend import JupyterBackend

        backend = JupyterBackend()
        assert backend.get_kernel_info("nonexistent") is None

    def test_get_kernel_connection_info_not_found(self):
        """get_kernel_connection_info should return None for unknown kernel."""
        from skop_runner.jupyter_backend import JupyterBackend

        backend = JupyterBackend()
        assert backend.get_kernel_connection_info("nonexistent") is None

    @pytest.mark.asyncio
    async def test_validate_kernel_not_found(self):
        """validate_kernel should return False for unknown kernel."""
        from skop_runner.jupyter_backend import JupyterBackend

        backend = JupyterBackend()
        assert await backend.validate_kernel("nonexistent") is False

    @pytest.mark.asyncio
    async def test_shutdown_all_empty(self):
        """shutdown_all should return empty results when no kernels."""
        from skop_runner.jupyter_backend import JupyterBackend

        backend = JupyterBackend()
        result = await backend.shutdown_all()
        assert result == {"stopped": [], "failed": {}}

    @pytest.mark.asyncio
    async def test_stop_kernel_not_found(self):
        """stop_kernel should raise KernelNotFoundError for unknown kernel."""
        from skop_runner.jupyter_backend import JupyterBackend

        backend = JupyterBackend()
        with pytest.raises(KernelNotFoundError):
            await backend.stop_kernel("nonexistent")

    @pytest.mark.asyncio
    async def test_restart_kernel_not_found(self):
        """restart_kernel should raise KernelNotFoundError for unknown kernel."""
        from skop_runner.jupyter_backend import JupyterBackend

        backend = JupyterBackend()
        with pytest.raises(KernelNotFoundError):
            await backend.restart_kernel("nonexistent")

    @pytest.mark.asyncio
    async def test_interrupt_kernel_not_found(self):
        """interrupt_kernel should raise KernelNotFoundError for unknown kernel."""
        from skop_runner.jupyter_backend import JupyterBackend

        backend = JupyterBackend()
        with pytest.raises(KernelNotFoundError):
            await backend.interrupt_kernel("nonexistent")

    def test_get_execution_count_not_found(self):
        """get_execution_count should raise KernelNotFoundError for unknown kernel."""
        from skop_runner.jupyter_backend import JupyterBackend

        backend = JupyterBackend()
        with pytest.raises(KernelNotFoundError):
            backend.get_execution_count("nonexistent")


class TestJupyterBackendCallbacks:
    """Test callback handling in JupyterBackend."""

    def test_callbacks_stored(self):
        """Callbacks should be stored in backend."""
        from skop_runner.jupyter_backend import JupyterBackend

        def on_died():
            pass

        def on_kernel_died(kid):
            pass

        def on_restarted():
            pass

        def on_status_change(kid, state):
            pass

        backend = JupyterBackend(
            on_backend_died=on_died,
            on_kernel_died=on_kernel_died,
            on_backend_restarted=on_restarted,
            on_kernel_status_change=on_status_change,
        )

        assert backend._on_backend_died is on_died
        assert backend._on_kernel_died is on_kernel_died
        assert backend._on_backend_restarted is on_restarted
        assert backend._on_kernel_status_change is on_status_change

    def test_handle_kernel_died_calls_callback(self):
        """_handle_kernel_died should invoke on_kernel_died with kernel_id."""
        from skop_runner.jupyter_backend import JupyterBackend

        received_ids = []
        backend = JupyterBackend(on_kernel_died=lambda kid: received_ids.append(kid))
        backend._handle_kernel_died("kernel-abc")

        assert received_ids == ["kernel-abc"]

    def test_handle_kernel_died_no_callback(self):
        """_handle_kernel_died should not error when callback is None."""
        from skop_runner.jupyter_backend import JupyterBackend

        backend = JupyterBackend()
        # Should not raise
        backend._handle_kernel_died("kernel-xyz")

    def test_handle_kernel_died_callback_exception_swallowed(self):
        """_handle_kernel_died should swallow callback exceptions."""
        from skop_runner.jupyter_backend import JupyterBackend

        def bad_callback(kid):
            raise RuntimeError("callback error")

        backend = JupyterBackend(on_kernel_died=bad_callback)
        # Should not raise
        backend._handle_kernel_died("kernel-123")

    def test_handle_kernel_status_change_calls_callback(self):
        """_handle_kernel_status_change should invoke callback with kernel_id and state."""
        from skop_runner.jupyter_backend import JupyterBackend
        from unittest.mock import MagicMock

        received = []
        backend = JupyterBackend(on_kernel_status_change=lambda kid, state: received.append((kid, state)))
        # Must have the kernel in _connections for callback to fire
        backend._connections["kernel-abc"] = MagicMock()

        backend._handle_kernel_status_change("kernel-abc", "busy")
        backend._handle_kernel_status_change("kernel-abc", "idle")

        assert received == [("kernel-abc", "busy"), ("kernel-abc", "idle")]

    def test_handle_kernel_status_change_ignores_unknown_kernel(self):
        """_handle_kernel_status_change should not fire for untracked kernels."""
        from skop_runner.jupyter_backend import JupyterBackend

        received = []
        backend = JupyterBackend(on_kernel_status_change=lambda kid, state: received.append((kid, state)))
        # No kernel in _connections
        backend._handle_kernel_status_change("unknown-kernel", "busy")

        assert received == []

    def test_handle_kernel_status_change_no_callback(self):
        """_handle_kernel_status_change should not error when callback is None."""
        from skop_runner.jupyter_backend import JupyterBackend
        from unittest.mock import MagicMock

        backend = JupyterBackend()
        backend._connections["kernel-abc"] = MagicMock()
        # Should not raise
        backend._handle_kernel_status_change("kernel-abc", "idle")

    def test_handle_kernel_status_change_exception_swallowed(self):
        """_handle_kernel_status_change should swallow callback exceptions."""
        from skop_runner.jupyter_backend import JupyterBackend
        from unittest.mock import MagicMock

        def bad_callback(kid, state):
            raise RuntimeError("callback error")

        backend = JupyterBackend(on_kernel_status_change=bad_callback)
        backend._connections["kernel-abc"] = MagicMock()
        # Should not raise
        backend._handle_kernel_status_change("kernel-abc", "busy")


class TestServerManagerCallbacks:
    """Test callback handling in JupyterServerManager."""

    def test_restart_callback_stored(self):
        """on_server_restarted should be stored in server manager."""
        from skop_runner.jupyter_backend.server_manager import JupyterServerManager

        def on_restarted():
            pass

        manager = JupyterServerManager(on_server_restarted=on_restarted)
        assert manager._on_server_restarted is on_restarted

    def test_restart_callback_default_none(self):
        """on_server_restarted should default to None."""
        from skop_runner.jupyter_backend.server_manager import JupyterServerManager

        manager = JupyterServerManager()
        assert manager._on_server_restarted is None

    @pytest.mark.asyncio
    async def test_initialize_wires_restart_callback(self):
        """initialize() should pass on_backend_restarted to server manager."""
        from skop_runner.jupyter_backend import JupyterBackend

        restarted_calls = []
        backend = JupyterBackend(on_backend_restarted=lambda: restarted_calls.append(True))
        await backend.initialize()

        assert backend._server_manager is not None
        assert backend._server_manager._on_server_restarted is not None

        # Invoke the server manager's callback to verify it bridges to backend
        backend._server_manager._on_server_restarted()
        assert restarted_calls == [True]
