"""
Jupyter Server backend implementing the KernelBackend protocol.

This module provides a complete implementation of the KernelBackend protocol
using Jupyter Server for kernel lifecycle management and code execution.
"""

import asyncio
import logging
import subprocess
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, AsyncIterator, Awaitable, Callable, Dict, List, Optional

from ..structured_log import get_logger as _get_slog

from ..kernel_backend import (
    KernelBackend,
    KernelBackendError,
    KernelConnectionError,
    KernelConnectionInfo,
    KernelGoneError,
    KernelInfo,
    KernelInputError,
    KernelNotFoundError,
    KernelStartError,
    KernelStatus,
    KernelTimeoutError,
    OnBackendDied,
    OnBackendRestarted,
    OnKernelDied,
    OnKernelStatusChange,
    register_backend,
)
from .api_client import (
    JupyterAPIClient,
    JupyterAPIError,
    JupyterConnectionError,
    JupyterGoneError,
    JupyterNotFoundError,
    JupyterServerError,
)
from .server_manager import JupyterServerManager
from .ws_client import (
    JupyterExecutionTimeoutError,
    JupyterKernelDiedError,
    JupyterKernelWebSocket,
    JupyterWebSocketConnectionError,
)

logger = logging.getLogger("skop-runner.jupyter_backend")
_log = _get_slog("jupyter_backend")

# Maximum execution time for a single cell (5 minutes)
DEFAULT_EXECUTION_TIMEOUT = 300.0


def _validate_python_path(python_path: str) -> None:
    """
    Validate that python_path is a valid Python executable.

    Security: Prevents execution of arbitrary binaries by verifying
    the path actually points to a Python interpreter.

    Raises:
        ValueError: If path is invalid or not a Python executable
    """
    path = Path(python_path)

    if not path.exists():
        raise ValueError(f"Python executable not found: {python_path}")

    if not path.is_file():
        raise ValueError(f"Python path is not a file: {python_path}")

    # Verify it's actually a Python executable
    try:
        result = subprocess.run(
            [python_path, "--version"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        output = result.stdout + result.stderr
        if result.returncode != 0 or "Python" not in output:
            raise ValueError(f"Not a valid Python executable: {python_path}")
    except subprocess.TimeoutExpired:
        raise ValueError(f"Python version check timed out: {python_path}")
    except FileNotFoundError:
        raise ValueError(f"Cannot execute Python path: {python_path}")
    except OSError as e:
        raise ValueError(f"Cannot verify Python executable: {e}")


@dataclass
class KernelConnection:
    """Holds connection info for a kernel managed via Jupyter Server."""

    kernel_id: str
    session_id: str
    websocket: JupyterKernelWebSocket
    python_path: Optional[str] = None
    env_name: Optional[str] = None
    execution_count: int = 0
    lock: asyncio.Lock = field(default_factory=asyncio.Lock)


class JupyterBackend:
    """
    Jupyter Server backend implementing the KernelBackend protocol.

    This backend uses Jupyter Server for kernel lifecycle management
    and WebSocket connections for code execution.

    Thread Safety:
        All async methods are safe to call concurrently from multiple
        coroutines. Per-kernel locking is used for execute_code().

    Attributes:
        _server_manager: JupyterServerManager for subprocess management
        _api_client: JupyterAPIClient for REST API calls
        _connections: Dict mapping kernel_id to KernelConnection
    """

    def __init__(
        self,
        on_backend_died: Optional[OnBackendDied] = None,
        on_kernel_died: Optional[OnKernelDied] = None,
        on_backend_restarted: Optional[OnBackendRestarted] = None,
        on_kernel_status_change: Optional[OnKernelStatusChange] = None,
        persistence_dir: Optional[str] = None,
    ):
        """
        Initialize the Jupyter backend.

        Args:
            on_backend_died: Callback when backend (Jupyter Server) dies
            on_kernel_died: Callback when a specific kernel dies (kernel_id passed)
            on_backend_restarted: Callback when backend restarts
            on_kernel_status_change: Callback when a kernel's execution state changes
                (called with kernel_id and execution_state e.g. "busy", "idle", "starting")
            persistence_dir: Optional directory for persisting server/session state to disk.
        """
        self._on_backend_died = on_backend_died
        self._on_kernel_died = on_kernel_died
        self._on_backend_restarted = on_backend_restarted
        self._on_kernel_status_change = on_kernel_status_change
        self._persistence_dir = persistence_dir

        self._server_manager: Optional[JupyterServerManager] = None
        self._api_client: Optional[JupyterAPIClient] = None
        self._connections: Dict[str, KernelConnection] = {}
        self._starting_kernel_ids: set[str] = set()
        self._status_task: Optional[asyncio.Task] = None
        self._initialized = False

    # =========================================================================
    # Protocol: Lifecycle Hooks
    # =========================================================================

    async def initialize(self, **kwargs: Any) -> None:
        """
        Initialize the backend.

        Args:
            root_dir: Optional initial root directory for Jupyter Server
        """
        if self._initialized:
            return

        root_dir = kwargs.get("root_dir")

        def handle_server_died() -> None:
            """Handle Jupyter Server death."""
            logger.error("Jupyter Server died")
            if self._on_backend_died:
                try:
                    self._on_backend_died()
                except Exception as e:
                    logger.error(f"on_backend_died callback failed: {e}")

        def handle_server_restarted() -> None:
            """Handle Jupyter Server restart."""
            logger.info("Jupyter Server restarted - existing connections may be stale")
            if self._on_backend_restarted:
                try:
                    self._on_backend_restarted()
                except Exception as e:
                    logger.error(f"on_backend_restarted callback failed: {e}")

        self._server_manager = JupyterServerManager(
            on_server_died=handle_server_died,
            on_server_restarted=handle_server_restarted,
            persistence_dir=self._persistence_dir,
        )

        if root_dir:
            await self._server_manager.ensure_running(root_dir=root_dir)
            self._api_client = JupyterAPIClient(
                self._server_manager.base_url,
                self._server_manager.token,
            )
            # Reconnect to any kernels surviving from previous runner session
            await self._reconnect_surviving_kernels()

        self._initialized = True
        self._status_task = asyncio.create_task(self._status_monitor())
        logger.info("JupyterBackend initialized")

    async def close(self, preserve_server: bool = False) -> None:
        """Clean up backend resources.

        Args:
            preserve_server: If True, leave Jupyter Server and kernels running
                for reuse on next runner startup. Only close our connections.
        """
        logger.info(f"Closing JupyterBackend (preserve_server={preserve_server})...")

        # Cancel status monitor
        if self._status_task and not self._status_task.done():
            self._status_task.cancel()
            try:
                await self._status_task
            except asyncio.CancelledError:
                pass
            self._status_task = None

        if preserve_server:
            # Soft close: disconnect WebSocket connections, leave kernels alive
            for conn in list(self._connections.values()):
                try:
                    await conn.websocket.close()
                except Exception:
                    pass
            self._connections.clear()
        else:
            # Hard close: kill all kernels
            await self.shutdown_all()

        # Close API client
        if self._api_client is not None:
            await self._api_client.close()
            self._api_client = None

        if preserve_server:
            # Detach from server without stopping it
            if self._server_manager is not None:
                await self._server_manager.detach()
                self._server_manager = None
        else:
            # Stop Jupyter Server
            if self._server_manager is not None:
                await self._server_manager.stop_server()
                self._server_manager = None

        self._initialized = False
        logger.info("JupyterBackend closed")

    # =========================================================================
    # Protocol: Backend State
    # =========================================================================

    def is_ready(self) -> bool:
        """Check if backend is ready to accept operations."""
        return (
            self._initialized
            and self._server_manager is not None
            and self._server_manager.is_running()
        )

    async def set_root_directory(self, path: str) -> str:
        """
        Set the working directory for the backend.

        This will restart Jupyter Server if the root directory changes.

        Args:
            path: Absolute path to the new root directory

        Returns:
            Resolved absolute path that was set
        """
        resolved_path = str(Path(path).expanduser().resolve())
        await self._ensure_server_running(root_dir=resolved_path)
        await self._server_manager.set_root(resolved_path)

        # Update API client with new credentials if server restarted
        if self._api_client is not None:
            await self._api_client.update_credentials(
                self._server_manager.base_url,
                self._server_manager.token,
            )
        else:
            self._api_client = JupyterAPIClient(
                self._server_manager.base_url,
                self._server_manager.token,
            )

        # Reconnect to surviving kernels before pruning — idempotent,
        # skips already-tracked kernels.  This is the "first time we have
        # a running server" moment after a runner restart.
        await self._reconnect_surviving_kernels()

        # Now prune anything still untracked (truly orphaned kernels)
        await self._prune_dead_connections()

        return resolved_path

    # =========================================================================
    # Protocol: Kernel Lifecycle
    # =========================================================================

    async def start_kernel(
        self,
        kernel_name: str = "python3",
        python_path: Optional[str] = None,
        env_name: Optional[str] = None,
        cwd: Optional[str] = None,
    ) -> str:
        """
        Start a new kernel and return its ID.

        Args:
            kernel_name: Kernel specification name (default: "python3")
            python_path: Path to Python executable for custom environments
            env_name: Display name for the environment
            cwd: Working directory for the kernel process

        Returns:
            Unique kernel ID string

        Raises:
            KernelStartError: If kernel fails to start
        """
        _t0 = time.perf_counter()
        try:
            # Ensure Jupyter Server is running
            await self._ensure_server_running(root_dir=cwd)

            # Prune connections whose kernels no longer exist in Jupyter Server
            await self._prune_dead_connections()

            # Handle custom Python path
            actual_kernel_name = kernel_name
            if python_path:
                _validate_python_path(python_path)
                # Register ephemeral kernelspec for custom Python
                spec_name = await self._server_manager.register_ephemeral_kernelspec(
                    python_path, env_name
                )
                actual_kernel_name = spec_name
                logger.info(f"Registered ephemeral kernelspec: {spec_name}")

            # Create session (which starts the kernel)
            session_path = f"session-{uuid.uuid4().hex[:8]}.ipynb"
            logger.info(f"Starting kernel with spec: {actual_kernel_name}")

            session = await self._api_client.create_session(
                path=session_path,
                kernel_name=actual_kernel_name,
            )

            kernel_id = session["kernel"]["id"]
            session_id = session["id"]

            # Mark as in-progress so _prune_dead_connections doesn't kill it
            # during the window between create_session and storing in _connections
            self._starting_kernel_ids.add(kernel_id)

            # Connect WebSocket for execution
            # If this fails, we must clean up the orphaned session/kernel
            try:
                ws = JupyterKernelWebSocket(
                    kernel_id=kernel_id,
                    base_url=self._server_manager.base_url,
                    token=self._server_manager.token,
                    on_kernel_died=lambda: self._handle_kernel_died(kernel_id),
                    is_kernel_alive=self._make_kernel_health_checker(kernel_id),
                    on_status_change=lambda state: self._handle_kernel_status_change(kernel_id, state),
                )
                await ws.connect()
            except Exception as ws_error:
                # Clean up orphaned session/kernel in Jupyter Server
                self._starting_kernel_ids.discard(kernel_id)
                logger.warning(
                    f"WebSocket connection failed for kernel {kernel_id}, "
                    f"cleaning up orphaned session {session_id}"
                )
                try:
                    await self._api_client.delete_session(session_id)
                except Exception as cleanup_error:
                    logger.error(f"Failed to cleanup orphaned session: {cleanup_error}")
                    # Also try direct kernel deletion as fallback
                    try:
                        await self._api_client.delete_kernel(kernel_id)
                    except Exception:
                        pass
                raise ws_error

            # Store connection and clear in-progress flag
            self._connections[kernel_id] = KernelConnection(
                kernel_id=kernel_id,
                session_id=session_id,
                websocket=ws,
                python_path=python_path,
                env_name=env_name,
            )
            self._starting_kernel_ids.discard(kernel_id)

            _log.info(
                "kernel.start",
                kernel_id=kernel_id,
                spec=actual_kernel_name,
                duration_ms=round((time.perf_counter() - _t0) * 1000),
                active_kernels=len(self._connections),
            )
            return kernel_id

        except JupyterConnectionError as e:
            raise KernelStartError(f"Failed to connect to Jupyter Server: {e}") from e
        except JupyterAPIError as e:
            raise KernelStartError(f"Failed to start kernel: {e}") from e
        except JupyterWebSocketConnectionError as e:
            raise KernelStartError(f"Failed to connect to kernel WebSocket: {e}") from e
        except ValueError as e:
            # Python path validation error
            raise KernelStartError(str(e)) from e

    async def stop_kernel(self, kernel_id: str) -> None:
        """
        Stop a running kernel.

        Args:
            kernel_id: ID of kernel to stop

        Raises:
            KernelNotFoundError: If kernel doesn't exist
        """
        _t0 = time.perf_counter()
        conn = self._get_connection(kernel_id)

        # Interrupt first to unblock any pending execution holding conn.lock
        # (e.g., cell waiting for input()). Without this, stop would deadlock
        # waiting for the lock that execute_code holds indefinitely.
        try:
            await self._api_client.interrupt_kernel(kernel_id)
        except Exception:
            pass  # Best-effort — kernel may already be dead

        # Try to acquire lock with timeout — force-close WebSocket if stuck
        lock_acquired = False
        try:
            await asyncio.wait_for(conn.lock.acquire(), timeout=5.0)
            lock_acquired = True
        except asyncio.TimeoutError:
            logger.warning(f"Lock timeout stopping kernel {kernel_id}, force-closing WebSocket")
            await conn.websocket.close()
            # Stuck execution should now get ConnectionClosed and release lock
            await asyncio.sleep(0.1)  # Let event loop process the close
            try:
                await asyncio.wait_for(conn.lock.acquire(), timeout=5.0)
                lock_acquired = True
            except asyncio.TimeoutError:
                logger.error(f"Failed to acquire lock for kernel {kernel_id}, forcing cleanup")
                self._connections.pop(kernel_id, None)
                logger.info(
                    f"Kernel {kernel_id} force-removed "
                    f"[active: {len(self._connections)}, "
                    f"ids: {list(self._connections.keys())}]"
                )
                return

        try:
            logger.info(f"Stopping kernel {kernel_id}")

            # Close WebSocket connection
            await conn.websocket.close()

            # Delete session (which stops the kernel)
            try:
                await self._api_client.delete_session(conn.session_id)
            except JupyterNotFoundError:
                # Session already gone
                pass
            except JupyterGoneError:
                # Kernel already dead
                pass
            except Exception as e:
                logger.warning(f"Error deleting session: {e}")
                # Try to delete kernel directly as fallback
                try:
                    await self._api_client.delete_kernel(kernel_id)
                except Exception:
                    pass

            # Idempotent: _on_kernel_died may have already popped this
            self._connections.pop(kernel_id, None)
            _log.info(
                "kernel.stop",
                kernel_id=kernel_id,
                duration_ms=round((time.perf_counter() - _t0) * 1000),
                active_kernels=len(self._connections),
            )
        finally:
            if lock_acquired:
                conn.lock.release()

    async def restart_kernel(self, kernel_id: str) -> None:
        """
        Restart a kernel, clearing its state.

        Args:
            kernel_id: ID of kernel to restart

        Raises:
            KernelNotFoundError: If kernel doesn't exist
            KernelStartError: If restart fails
        """
        _t0 = time.perf_counter()
        conn = self._get_connection(kernel_id)

        # Interrupt first to unblock any pending execution holding conn.lock
        # (e.g., cell waiting for input()). Without this, restart would deadlock
        # waiting for the lock that execute_code holds indefinitely.
        try:
            await self._api_client.interrupt_kernel(kernel_id)
        except Exception:
            pass  # Best-effort — kernel may already be dead

        # Try to acquire lock with timeout — force-close WebSocket if stuck
        lock_acquired = False
        try:
            await asyncio.wait_for(conn.lock.acquire(), timeout=5.0)
            lock_acquired = True
        except asyncio.TimeoutError:
            logger.warning(f"Lock timeout restarting kernel {kernel_id}, force-closing WebSocket")
            await conn.websocket.close()
            await asyncio.sleep(0.1)
            try:
                await asyncio.wait_for(conn.lock.acquire(), timeout=5.0)
                lock_acquired = True
            except asyncio.TimeoutError:
                logger.error(f"Failed to acquire lock for kernel {kernel_id}, forcing cleanup")
                self._connections.pop(kernel_id, None)
                raise KernelConnectionError(
                    f"Kernel {kernel_id} is stuck and could not be restarted. "
                    "Please stop the kernel and start a new one."
                )

        try:
            logger.info(f"Restarting kernel {kernel_id}")

            try:
                # Close existing WebSocket
                await conn.websocket.close()

                # Restart via API
                await self._api_client.restart_kernel(kernel_id)

                # Reconnect WebSocket
                conn.websocket = JupyterKernelWebSocket(
                    kernel_id=kernel_id,
                    base_url=self._server_manager.base_url,
                    token=self._server_manager.token,
                    on_kernel_died=lambda: self._handle_kernel_died(kernel_id),
                    is_kernel_alive=self._make_kernel_health_checker(kernel_id),
                    on_status_change=lambda state: self._handle_kernel_status_change(kernel_id, state),
                )
                await conn.websocket.connect()

                # Reset execution count
                conn.execution_count = 0

                # Ensure conn is tracked — _on_kernel_died may have popped it
                # during the force-close path
                self._connections[kernel_id] = conn

                _log.info(
                    "kernel.restart",
                    kernel_id=kernel_id,
                    duration_ms=round((time.perf_counter() - _t0) * 1000),
                )

            except (JupyterNotFoundError, JupyterGoneError) as e:
                raise KernelNotFoundError(f"Kernel {kernel_id} not found") from e
            except JupyterAPIError as e:
                raise KernelStartError(f"Failed to restart kernel: {e}") from e
            except JupyterWebSocketConnectionError as e:
                raise KernelConnectionError(
                    f"Failed to reconnect to kernel: {e}"
                ) from e
        finally:
            if lock_acquired:
                conn.lock.release()

    async def interrupt_kernel(self, kernel_id: str) -> None:
        """
        Send interrupt signal to kernel.

        Args:
            kernel_id: ID of kernel to interrupt

        Raises:
            KernelNotFoundError: If kernel doesn't exist
        """
        # Verify kernel exists
        self._get_connection(kernel_id)

        try:
            logger.info(f"Interrupting kernel {kernel_id}")
            await self._api_client.interrupt_kernel(kernel_id)
        except JupyterNotFoundError as e:
            raise KernelNotFoundError(f"Kernel {kernel_id} not found") from e

    async def shutdown_all(self, timeout: float = 15.0) -> Dict[str, Any]:
        """
        Shutdown all running kernels.

        Args:
            timeout: Maximum time to wait for each kernel shutdown

        Returns:
            Dict with 'stopped' list and 'failed' dict
        """
        kernel_ids = list(self._connections.keys())
        if not kernel_ids:
            return {"stopped": [], "failed": {}}

        stopped = []
        failed = {}

        async def stop_one(kernel_id: str):
            try:
                await asyncio.wait_for(self.stop_kernel(kernel_id), timeout=timeout)
                return (kernel_id, None)
            except asyncio.TimeoutError:
                logger.warning(f"Kernel {kernel_id} shutdown timed out")
                # Force cleanup
                conn = self._connections.pop(kernel_id, None)
                if conn:
                    try:
                        await conn.websocket.close()
                    except Exception:
                        pass
                return (kernel_id, None)
            except Exception as e:
                return (kernel_id, str(e))

        results = await asyncio.gather(*[stop_one(kid) for kid in kernel_ids])

        for kernel_id, error in results:
            if error:
                failed[kernel_id] = error
                logger.error(f"Error stopping kernel {kernel_id}: {error}")
            else:
                stopped.append(kernel_id)

        if failed:
            logger.warning(
                f"Failed to stop {len(failed)} kernel(s): {list(failed.keys())}"
            )

        return {"stopped": stopped, "failed": failed}

    # =========================================================================
    # Protocol: Code Execution
    # =========================================================================

    async def execute_code(
        self,
        kernel_id: str,
        code: str,
        timeout: float = DEFAULT_EXECUTION_TIMEOUT,
    ) -> AsyncIterator[dict]:
        """
        Execute code and stream outputs.

        Args:
            kernel_id: ID of kernel to execute in
            code: Code string to execute
            timeout: Maximum execution time in seconds

        Yields:
            Output dicts with 'output_type' key

        Raises:
            KernelNotFoundError: If kernel doesn't exist
            KernelTimeoutError: If execution times out
            KernelConnectionError: If connection to kernel is lost
        """
        conn = self._get_connection(kernel_id)

        async with conn.lock:
            conn.execution_count += 1

            logger.info(
                f"Executing code in kernel {kernel_id} (exec #{conn.execution_count})"
            )

            _t0 = time.perf_counter()
            _output_size = 0
            try:
                async for output in conn.websocket.execute(code, timeout=timeout):
                    _output_size += len(str(output))
                    yield output
            except JupyterExecutionTimeoutError as e:
                raise KernelTimeoutError(str(e)) from e
            except JupyterWebSocketConnectionError as e:
                raise KernelConnectionError(str(e)) from e
            except JupyterKernelDiedError as e:
                raise KernelGoneError(str(e)) from e
            except asyncio.CancelledError:
                # Task was cancelled - interrupt the kernel
                logger.warning(
                    f"Execution cancelled for kernel {kernel_id}, interrupting..."
                )
                try:
                    await self._api_client.interrupt_kernel(kernel_id)
                except Exception as e:
                    logger.warning(f"Failed to interrupt kernel on cancellation: {e}")
                raise

            _log.info(
                "kernel.execute",
                kernel_id=kernel_id,
                code_length=len(code),
                output_size=_output_size,
                duration_ms=round((time.perf_counter() - _t0) * 1000),
            )

    async def send_input_reply(self, kernel_id: str, value: str) -> None:
        """
        Send stdin input to a kernel waiting for input().

        Args:
            kernel_id: ID of kernel waiting for input
            value: Input value to send

        Raises:
            KernelNotFoundError: If kernel doesn't exist
            KernelInputError: If kernel is not waiting for input
        """
        conn = self._get_connection(kernel_id)

        if not conn.websocket.awaiting_input:
            raise KernelInputError("Kernel is not waiting for input")

        try:
            await conn.websocket.send_input(value)
        except ValueError as e:
            raise KernelInputError(str(e)) from e

    # =========================================================================
    # Protocol: Status and Discovery
    # =========================================================================

    async def get_kernel_status(self, kernel_id: str) -> Dict[str, Any]:
        """
        Get kernel status information.

        Args:
            kernel_id: ID of kernel to query

        Returns:
            Dict with status, kernel_id, and execution_state
        """
        if kernel_id not in self._connections:
            return {"status": "not_found", "kernel_id": kernel_id}

        try:
            kernel_info = await self._api_client.get_kernel(kernel_id)
            execution_state = kernel_info.get("execution_state", "unknown")

            # Map Jupyter states to our status
            if execution_state in ("idle", "busy", "starting"):
                status = "running"
            else:
                status = execution_state

            return {
                "status": status,
                "kernel_id": kernel_id,
                "execution_state": execution_state,
            }
        except Exception as e:
            logger.warning(f"Error getting kernel status: {e}")
            return {"status": "unknown", "kernel_id": kernel_id}

    def get_kernel_info(self, kernel_id: str) -> Optional[KernelInfo]:
        """
        Get kernel metadata (synchronous).

        Args:
            kernel_id: ID of kernel to query

        Returns:
            KernelInfo if kernel exists, None otherwise
        """
        conn = self._connections.get(kernel_id)
        if conn is None:
            return None

        return KernelInfo(
            kernel_id=conn.kernel_id,
            status=KernelStatus.UNKNOWN,  # Would need async call for accurate status
            python_path=conn.python_path,
            env_name=conn.env_name,
            execution_count=conn.execution_count,
            metadata={"session_id": conn.session_id},
        )

    def get_kernel_connection_info(self, kernel_id: str) -> Optional[KernelConnectionInfo]:
        """
        Get connection info for a kernel.

        Used by SessionManager to track kernel-to-notebook associations.

        Args:
            kernel_id: ID of kernel to query

        Returns:
            KernelConnectionInfo if kernel exists, None otherwise
        """
        conn = self._connections.get(kernel_id)
        if conn is None:
            return None

        return KernelConnectionInfo(
            session_id=conn.session_id,
            python_path=conn.python_path,
            env_name=conn.env_name,
        )

    def list_kernels(self) -> List[Dict[str, Any]]:
        """
        List all running kernels (synchronous).

        Returns:
            List of dicts with kernel_id, status, python_path, env_name
        """
        return [
            {
                "kernel_id": kid,
                "status": "running",
                "python_path": conn.python_path,
                "env_name": conn.env_name,
            }
            for kid, conn in self._connections.items()
        ]

    async def list_available_kernels(self) -> List[Dict[str, Any]]:
        """
        List available kernel specifications.

        Returns:
            List of dicts with name, display_name, language
        """
        await self._ensure_server_running()

        specs = await self._api_client.get_kernelspecs()
        kernelspecs = specs.get("kernelspecs", {})

        return [
            {
                "name": name,
                "display_name": spec.get("spec", {}).get("display_name", name),
                "language": spec.get("spec", {}).get("language", "unknown"),
            }
            for name, spec in kernelspecs.items()
        ]

    def get_execution_count(self, kernel_id: str) -> int:
        """
        Get the current execution count for a kernel.

        Args:
            kernel_id: ID of kernel to query

        Returns:
            Current execution count

        Raises:
            KernelNotFoundError: If kernel doesn't exist
        """
        conn = self._get_connection(kernel_id)
        return conn.execution_count

    async def validate_kernel(self, kernel_id: str) -> bool:
        """
        Check if kernel is still valid/running.

        Args:
            kernel_id: ID of kernel to validate

        Returns:
            True if kernel is valid and responsive
        """
        if kernel_id not in self._connections:
            return False

        try:
            status = await self.get_kernel_status(kernel_id)
            return status.get("status") in ("running", "idle", "busy")
        except Exception:
            return False

    # =========================================================================
    # Internal Methods
    # =========================================================================

    async def _status_monitor(self) -> None:
        """Periodically log kernel count and root directory."""
        try:
            while True:
                await asyncio.sleep(60)
                if self._connections:
                    root = self._server_manager._root_dir if self._server_manager else None
                    logger.info(
                        f"Active kernels: {len(self._connections)}"
                        f" | root: {root}"
                    )
        except asyncio.CancelledError:
            pass

    async def _ensure_server_running(self, root_dir: Optional[str] = None) -> None:
        """
        Ensure the Jupyter Server is running and API client is ready.

        Args:
            root_dir: Project root directory for the server
        """
        if self._server_manager is None:
            raise KernelBackendError("JupyterBackend not initialized")

        await self._server_manager.ensure_running(root_dir=root_dir)

        # Create or update API client if needed
        if self._api_client is None or self._api_client.base_url != self._server_manager.base_url:
            if self._api_client is not None:
                await self._api_client.close()
            self._api_client = JupyterAPIClient(
                self._server_manager.base_url,
                self._server_manager.token,
            )

    def _make_kernel_health_checker(self, kernel_id: str) -> Callable[[], Awaitable[bool]]:
        """Create an async health check callback for a kernel."""
        async def check() -> bool:
            try:
                info = await self._api_client.get_kernel(kernel_id)
                state = info.get("execution_state", "unknown")
                # Only "busy" means our execution is still running.
                # "idle" after 10s of silence = execution lost (kernel restarted).
                # "dead"/"starting"/"restarting" = kernel died or is restarting.
                return state == "busy"
            except (JupyterNotFoundError, JupyterGoneError, JupyterServerError):
                return False
            except Exception:
                return True  # Assume alive on transient error
        return check

    async def _reconnect_surviving_kernels(self) -> int:
        """
        Reconnect WebSocket to kernels surviving from a previous runner session.

        Queries Jupyter Server for running kernels and creates WebSocket
        connections for any that aren't already tracked in ``_connections``.

        Returns:
            Number of kernels successfully reconnected.
        """
        if not self._api_client:
            return 0

        try:
            server_kernels = await self._api_client.list_kernels()
        except Exception as e:
            logger.warning(f"Failed to list kernels for reconnection: {e}")
            return 0

        tracked_ids = set(self._connections.keys())
        untracked = [k for k in server_kernels if k["id"] not in tracked_ids]
        if not untracked:
            return 0

        # Fetch sessions once for all kernels (avoid N API calls)
        try:
            all_sessions = await self._api_client.list_sessions()
        except Exception as e:
            logger.warning(f"Failed to list sessions for reconnection: {e}")
            return 0

        # Build kernel_id -> session_id lookup
        kernel_session_map: dict[str, str] = {}
        for sess in all_sessions:
            kid = sess.get("kernel", {}).get("id")
            if kid:
                kernel_session_map[kid] = sess["id"]

        reconnected = 0

        for kernel_info in untracked:
            kernel_id = kernel_info["id"]

            # This kernel exists in Jupyter but we don't have a connection
            _log.info("kernel.reconnect_attempt", kernel_id=kernel_id)
            logger.info(f"Reconnecting to surviving kernel {kernel_id}...")

            try:
                session_id = kernel_session_map.get(kernel_id)

                if session_id is None:
                    # Kernel has no session — it's orphaned, clean it up
                    logger.info(f"Kernel {kernel_id} has no session, stopping it")
                    try:
                        await self._api_client.delete_kernel(kernel_id)
                    except Exception:
                        pass
                    continue

                # Create WebSocket connection
                ws = JupyterKernelWebSocket(
                    kernel_id=kernel_id,
                    base_url=self._server_manager.base_url,
                    token=self._server_manager.token,
                    on_kernel_died=lambda kid=kernel_id: self._handle_kernel_died(kid),
                    is_kernel_alive=self._make_kernel_health_checker(kernel_id),
                    on_status_change=lambda state, kid=kernel_id: self._handle_kernel_status_change(kid, state),
                )
                await ws.connect()

                self._connections[kernel_id] = KernelConnection(
                    kernel_id=kernel_id,
                    session_id=session_id,
                    websocket=ws,
                )
                reconnected += 1
                _log.info("kernel.reconnect_success", kernel_id=kernel_id)
                logger.info(
                    f"Reconnected to kernel {kernel_id} (session {session_id}) "
                    f"[active: {len(self._connections)}]"
                )

            except Exception as e:
                _log.warning("kernel.reconnect_failed", kernel_id=kernel_id, error=str(e))
                logger.warning(
                    f"Failed to reconnect to kernel {kernel_id}: {e}, stopping it"
                )
                # Clean up the orphaned kernel
                try:
                    await self._api_client.delete_kernel(kernel_id)
                except Exception:
                    pass

        if reconnected:
            logger.info(f"Reconnected to {reconnected} surviving kernel(s)")
        return reconnected

    async def _prune_dead_connections(self) -> None:
        """Remove stale connections AND stop untracked kernels."""
        if not self._api_client:
            return

        try:
            server_kernels = await self._api_client.list_kernels()
            alive_ids = {k["id"] for k in server_kernels}
        except Exception as e:
            logger.warning(f"Failed to list kernels for pruning: {e}")
            return

        # Direction 1: Remove _connections entries for dead kernels
        dead_ids = [kid for kid in self._connections if kid not in alive_ids]
        for kid in dead_ids:
            conn = self._connections.pop(kid, None)
            if conn:
                try:
                    await conn.websocket.close()
                except Exception:
                    pass
            logger.info(
                f"Pruned orphaned connection: {kid} "
                f"[active: {len(self._connections)}, "
                f"ids: {list(self._connections.keys())}]"
            )

        # Direction 2: Stop Jupyter kernels not in _connections
        # Skip kernels that are currently being started (between create_session and
        # storing in _connections) to avoid killing a concurrent start_kernel's kernel
        tracked_ids = set(self._connections.keys())
        untracked = [
            k for k in server_kernels
            if k["id"] not in tracked_ids
            and k["id"] not in self._starting_kernel_ids
        ]
        for kernel in untracked:
            kid = kernel["id"]
            logger.info(f"Stopping untracked kernel: {kid}")
            try:
                await self._api_client.delete_kernel(kid)
            except Exception:
                pass


    def _get_connection(self, kernel_id: str) -> KernelConnection:
        """Get kernel connection or raise error."""
        if kernel_id not in self._connections:
            raise KernelNotFoundError(f"Kernel {kernel_id} not found")
        return self._connections[kernel_id]

    def _handle_kernel_status_change(self, kernel_id: str, execution_state: str) -> None:
        """Handle kernel status change from WebSocket IOPub channel."""
        if kernel_id not in self._connections:
            return
        if self._on_kernel_status_change:
            try:
                self._on_kernel_status_change(kernel_id, execution_state)
            except Exception as e:
                logger.warning(f"on_kernel_status_change callback failed: {e}")

    def _handle_kernel_died(self, kernel_id: str) -> None:
        """Handle kernel death notification from WebSocket."""
        _log.warning("kernel.died", kernel_id=kernel_id)

        # Remove dead connection so it doesn't block restart/start
        conn = self._connections.pop(kernel_id, None)
        if conn:
            logger.warning(
                f"Kernel {kernel_id} removed (died) "
                f"[active: {len(self._connections)}, "
                f"ids: {list(self._connections.keys())}]"
            )
            try:
                # Fire-and-forget close; the WebSocket is likely already dead
                asyncio.get_running_loop().create_task(conn.websocket.close())
            except Exception:
                pass

        if self._on_kernel_died:
            try:
                self._on_kernel_died(kernel_id)
            except Exception as e:
                logger.error(f"on_kernel_died callback failed: {e}")

    # =========================================================================
    # Optional Capabilities (not part of KernelBackend protocol)
    # =========================================================================

    async def get_kernelspecs(self) -> Optional[dict]:
        """
        Get raw kernelspecs from Jupyter Server.

        This is an optional capability not part of the KernelBackend protocol.
        Used by env_manager to discover non-Python kernels (R, Julia, etc.).

        Returns:
            Raw kernelspecs dict from Jupyter Server, or None if unavailable.
        """
        if self._api_client is None:
            return None
        try:
            return await self._api_client.get_kernelspecs()
        except Exception as e:
            logger.debug(f"Could not get kernelspecs: {e}")
            return None


# Register the Jupyter backend
register_backend("jupyter", JupyterBackend)
