"""
Backend-agnostic kernel management protocol.

This module defines the interface that all kernel backends must implement,
allowing Jupyter Server to be swapped for alternative implementations
(e.g., direct ZMQ kernel management, custom process management).

Key components:
- KernelBackend: Protocol that all backends must implement
- Exception hierarchy: Backend-agnostic error types
- Data classes: KernelInfo, KernelStatus, AvailableEnvironment
- Backend registry: Factory pattern for backend instantiation
"""

import threading
from dataclasses import dataclass, field
from enum import Enum
from typing import (
    Any,
    AsyncIterator,
    Callable,
    Dict,
    List,
    Optional,
    Protocol,
    TypedDict,
    runtime_checkable,
)


# =============================================================================
# Exceptions (backend-agnostic)
# =============================================================================


class KernelBackendError(Exception):
    """Base exception for kernel backend errors.

    All backend implementations should translate their internal exceptions
    to appropriate subclasses of this exception.
    """
    pass


class KernelNotFoundError(KernelBackendError):
    """Kernel with given ID does not exist.

    Raised when attempting to operate on a kernel that doesn't exist
    or has already been stopped.
    """
    pass


class KernelStartError(KernelBackendError):
    """Failed to start kernel.

    Raised when kernel creation fails due to:
    - Invalid Python path
    - Missing dependencies (e.g., ipykernel not installed)
    - Resource limits exceeded
    - Backend not ready
    """
    pass


class KernelConnectionError(KernelBackendError):
    """Lost connection to kernel.

    Raised when:
    - WebSocket connection fails or is lost
    - Kernel becomes unreachable
    - Network issues prevent communication
    """
    pass


class KernelExecutionError(KernelBackendError):
    """Error during code execution.

    Note: This is for execution infrastructure errors, not Python exceptions
    raised by user code (those are returned as error outputs).
    """
    pass


class KernelTimeoutError(KernelBackendError):
    """Operation timed out.

    Raised when:
    - Code execution exceeds timeout
    - Kernel startup takes too long
    - API calls don't respond in time
    """
    pass


class KernelInputError(KernelBackendError):
    """Error with stdin input handling.

    Raised when:
    - Attempting to send input when kernel is not waiting
    - Input reply fails to deliver
    """
    pass


class KernelBusyError(KernelBackendError):
    """Kernel is already executing code.

    Raised when attempting to execute code on a kernel that is
    already busy and the backend doesn't support queueing.
    """
    pass


class KernelResourceConflictError(KernelBackendError):
    """Resource conflict (e.g., concurrent session operations).

    Raised when:
    - Concurrent operations conflict
    - Session already exists for notebook
    """
    pass


class KernelGoneError(KernelBackendError):
    """Kernel died unexpectedly.

    Raised when the kernel process terminates unexpectedly
    during an operation.
    """
    pass


# =============================================================================
# Data Classes (backend-agnostic)
# =============================================================================


class KernelStatus(str, Enum):
    """Possible kernel execution states.

    These map to Jupyter's kernel states but are backend-agnostic.
    """
    STARTING = "starting"    # Kernel is starting up
    IDLE = "idle"            # Kernel is ready for execution
    BUSY = "busy"            # Kernel is executing code
    RESTARTING = "restarting"  # Kernel is restarting
    DEAD = "dead"            # Kernel has terminated
    NOT_FOUND = "not_found"  # Kernel ID doesn't exist
    UNKNOWN = "unknown"      # Status cannot be determined


@dataclass
class KernelInfo:
    """Backend-agnostic kernel metadata.

    This provides the information that higher layers need about a kernel
    without exposing backend-specific implementation details.

    Attributes:
        kernel_id: Unique identifier for the kernel
        status: Current execution state
        python_path: Path to Python executable (if applicable)
        env_name: Display name of the environment
        execution_count: Number of executions performed
        metadata: Backend-specific metadata (opaque to consumers)
    """
    kernel_id: str
    status: KernelStatus = KernelStatus.UNKNOWN
    python_path: Optional[str] = None
    env_name: Optional[str] = None
    execution_count: int = 0
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __repr__(self) -> str:
        return (
            f"KernelInfo(id={self.kernel_id!r}, "
            f"status={self.status.value}, "
            f"env={self.env_name!r})"
        )


@dataclass
class AvailableEnvironment:
    """An environment that can be used to start a kernel.

    Represents a Python environment or kernel specification that
    is available for starting kernels.

    Attributes:
        name: Unique identifier for the environment
        display_name: Human-readable name (e.g., "Python 3.11 (.venv)")
        env_type: Type of environment ("venv", "conda", "pyenv", "system", "kernelspec:*")
        python_path: Path to Python executable (None for non-Python kernels)
        python_version: Version string (e.g., "3.11.5")
        has_ipykernel: Whether ipykernel is installed
        kernel_name: Kernelspec name (None for Python envs using ephemeral specs)
        language: Programming language ("python", "r", "julia", etc.)
        is_active: Whether this is the currently active environment
        path: Environment directory path (for deduplication)
    """
    name: str
    display_name: str
    env_type: str
    python_path: Optional[str] = None
    python_version: Optional[str] = None
    has_ipykernel: bool = False
    kernel_name: Optional[str] = None
    language: str = "python"
    is_active: bool = False
    path: Optional[str] = None

    def __repr__(self) -> str:
        return (
            f"AvailableEnvironment(name={self.name!r}, "
            f"type={self.env_type}, "
            f"python={self.python_version})"
        )


class KernelConnectionInfo(TypedDict):
    """Connection info returned by get_kernel_connection_info().

    Used by SessionManager to track kernel-to-notebook associations.
    """
    session_id: str
    python_path: Optional[str]
    env_name: Optional[str]


# =============================================================================
# Callback Types
# =============================================================================

# Called when the backend dies and cannot be restarted
OnBackendDied = Callable[[], None]

# Called when a specific kernel dies unexpectedly (kernel_id passed)
OnKernelDied = Callable[[str], None]

# Called when the backend restarts (connections may be stale)
OnBackendRestarted = Callable[[], None]

# Called when a kernel's execution state changes (kernel_id, execution_state)
OnKernelStatusChange = Callable[[str, str], None]


# =============================================================================
# Protocol Definition
# =============================================================================


@runtime_checkable
class KernelBackend(Protocol):
    """Protocol defining the kernel backend interface.

    All kernel backends (Jupyter Server, direct ZMQ, etc.) must implement
    this interface. The Protocol is @runtime_checkable for isinstance checks.

    Thread Safety:
        All async methods must be safe to call concurrently from multiple
        coroutines. Implementations should use appropriate locking.

    Concurrency Contract:
        - execute_code() calls on the SAME kernel are serialized internally
        - execute_code() calls on DIFFERENT kernels may run concurrently
        - Implementations MUST use per-kernel locking to enforce this
        - If a kernel is busy and queueing is not supported, raise KernelBusyError

    Error Handling:
        Methods should raise appropriate KernelBackendError subclasses.
        Implementations must not leak backend-specific exceptions.

    Output Format:
        execute_code() yields dicts following the Jupyter messaging protocol:
        - stream: {"output_type": "stream", "name": "stdout/stderr", "text": "..."}
        - execute_result: {"output_type": "execute_result", "data": {...}, ...}
        - display_data: {"output_type": "display_data", "data": {...}, ...}
        - error: {"output_type": "error", "ename": "...", "evalue": "...", "traceback": [...]}
        - input_request: {"output_type": "input_request", "prompt": "...", "password": bool}
    """

    # =========================================================================
    # 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.
                        If provided, an ephemeral kernelspec may be created.
            env_name: Display name for the environment (for UI purposes)
            cwd: Working directory for the kernel process

        Returns:
            Unique kernel ID string

        Raises:
            KernelStartError: If kernel fails to start (invalid path, missing
                            ipykernel, resource limits, etc.)
            KernelBackendError: If backend is not ready
        """
        ...

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

        Gracefully shuts down the kernel, closing connections and
        cleaning up resources.

        Args:
            kernel_id: ID of kernel to stop

        Raises:
            KernelNotFoundError: If kernel doesn't exist
        """
        ...

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

        The kernel ID remains the same after restart. All variables
        and imports are cleared. Execution count is reset.

        Args:
            kernel_id: ID of kernel to restart

        Raises:
            KernelNotFoundError: If kernel doesn't exist
            KernelStartError: If restart fails
        """
        ...

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

        Interrupts any running execution. The kernel remains running
        and can accept new executions.

        Args:
            kernel_id: ID of kernel to interrupt

        Raises:
            KernelNotFoundError: If kernel doesn't exist
        """
        ...

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

        Attempts graceful shutdown of all kernels in parallel.
        Kernels that don't respond within timeout are force-killed.

        Args:
            timeout: Maximum time to wait for each kernel shutdown (seconds)

        Returns:
            Dict with:
            - 'stopped': List of successfully stopped kernel IDs
            - 'failed': Dict mapping kernel IDs to error messages
        """
        ...

    # =========================================================================
    # Code Execution
    # =========================================================================

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

        Executes the given code in the kernel and yields output messages
        as they arrive. This is an async generator.

        Args:
            kernel_id: ID of kernel to execute in
            code: Code string to execute
            timeout: Maximum execution time in seconds (default: 5 minutes).
                    Use 0 for no timeout.

        Yields:
            Output dicts with 'output_type' key. See class docstring for format.

        Raises:
            KernelNotFoundError: If kernel doesn't exist
            KernelTimeoutError: If execution times out
            KernelConnectionError: If connection to kernel is lost
            KernelBusyError: If kernel is busy and queueing not supported

        Note:
            Executions on the same kernel are serialized (one at a time).
            Executions on different kernels can run concurrently.
        """
        ...

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

        When a kernel executes code that calls input(), it sends an
        input_request output. This method provides the response.

        Args:
            kernel_id: ID of kernel waiting for input
            value: Input value to send (the text the user typed)

        Raises:
            KernelNotFoundError: If kernel doesn't exist
            KernelInputError: If kernel is not waiting for input
        """
        ...

    # =========================================================================
    # Status and Discovery
    # =========================================================================

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

        Queries the kernel for its current status. This may involve
        network I/O to the kernel or backend server.

        Args:
            kernel_id: ID of kernel to query

        Returns:
            Dict with at least:
            - 'status': "running", "not_found", or "unknown"
            - 'kernel_id': The kernel ID
            - 'execution_state': "idle", "busy", "starting" (optional)
        """
        ...

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

        Returns cached information about the kernel without making
        network calls. For live status, use get_kernel_status().

        Args:
            kernel_id: ID of kernel to query

        Returns:
            KernelInfo if kernel exists, None otherwise
        """
        ...

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

        Returns information needed by SessionManager to track the
        kernel-to-notebook association.

        Implementations MUST return a non-None value for any kernel_id
        returned by a successful start_kernel() call. Returning None
        for a valid kernel_id will cause SessionManager to treat the
        kernel as orphaned and attempt cleanup.

        Args:
            kernel_id: ID of kernel to query

        Returns:
            KernelConnectionInfo if kernel exists, None otherwise
        """
        ...

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

        Returns cached information about all kernels without making
        network calls.

        Returns:
            List of dicts with kernel_id, status, python_path, env_name
        """
        ...

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

        Returns all kernels/environments that can be started. This
        may include Python environments and non-Python kernelspecs.

        Returns:
            List of dicts with name, display_name, language, and
            environment-specific fields
        """
        ...

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

        Returns the number of executions performed since the kernel
        was started (or last restarted).

        Args:
            kernel_id: ID of kernel to query

        Returns:
            Current execution count

        Raises:
            KernelNotFoundError: If kernel doesn't exist
        """
        ...

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

        Verifies that the kernel is still responsive. Used to detect
        stale connections after backend restarts.

        Args:
            kernel_id: ID of kernel to validate

        Returns:
            True if kernel is valid and responsive, False otherwise
        """
        ...

    # =========================================================================
    # Backend State
    # =========================================================================

    def is_ready(self) -> bool:
        """Check if backend is ready to accept operations.

        Returns True if the backend is initialized and can accept
        kernel operations. Returns False during startup or after
        fatal errors.

        Returns:
            True if ready, False otherwise
        """
        ...

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

        Changes the root directory for kernel operations. May trigger
        a backend restart if the backend doesn't support hot-swapping.

        Args:
            path: Absolute path to the new root directory

        Returns:
            Resolved absolute path that was set

        Raises:
            KernelBackendError: If the path is invalid or change fails
        """
        ...

    # =========================================================================
    # Lifecycle Hooks
    # =========================================================================

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

        Called once at startup. Implementations may start background
        processes, establish connections, etc.

        Args:
            **kwargs: Backend-specific initialization options
        """
        ...

    async def close(self) -> None:
        """Clean up backend resources.

        Called at shutdown. Implementations should stop all kernels,
        close connections, and release resources.
        """
        ...


# =============================================================================
# Backend Registry and Factory
# =============================================================================

_backend_registry: Dict[str, type] = {}
_registry_lock = threading.Lock()


def register_backend(name: str, backend_class: type) -> None:
    """Register a kernel backend implementation.

    Registers a backend class that can be instantiated via get_backend().
    Registration is thread-safe and allows overwriting existing entries.

    Args:
        name: Unique identifier for the backend (e.g., "jupyter", "zmq", "mock")
        backend_class: Class implementing KernelBackend protocol

    Raises:
        TypeError: If backend_class is not a class

    Example:
        >>> register_backend("jupyter", JupyterBackend)
        >>> register_backend("mock", MockKernelBackend)
    """
    if not isinstance(backend_class, type):
        raise TypeError(f"backend_class must be a class, got {type(backend_class)}")

    with _registry_lock:
        _backend_registry[name] = backend_class


def get_backend(name: str, **kwargs: Any) -> KernelBackend:
    """Get an instance of a registered backend.

    Creates and returns an instance of the specified backend.
    The backend is not initialized - call initialize() separately.

    Args:
        name: Backend identifier (e.g., "jupyter", "mock")
        **kwargs: Arguments passed to the backend constructor

    Returns:
        Instance of the requested backend

    Raises:
        ValueError: If backend name is not registered

    Example:
        >>> backend = get_backend("jupyter", on_backend_died=my_callback)
        >>> await backend.initialize()
    """
    with _registry_lock:
        if name not in _backend_registry:
            available = list(_backend_registry.keys())
            raise ValueError(
                f"Unknown backend '{name}'. "
                f"Available backends: {available if available else '(none registered)'}"
            )
        backend_class = _backend_registry[name]

    return backend_class(**kwargs)


def list_backends() -> List[str]:
    """List all registered backend names.

    Returns:
        List of registered backend identifiers

    Example:
        >>> list_backends()
        ['jupyter', 'mock']
    """
    with _registry_lock:
        return list(_backend_registry.keys())


def unregister_backend(name: str) -> bool:
    """Unregister a backend (primarily for testing).

    Args:
        name: Backend identifier to unregister

    Returns:
        True if backend was unregistered, False if it wasn't registered
    """
    with _registry_lock:
        if name in _backend_registry:
            del _backend_registry[name]
            return True
        return False
