"""
Jupyter Kernel WebSocket Client.

Provides async WebSocket client for Jupyter kernel communication with:
- Background reader for continuous message processing
- Per-execution asyncio.Queue routing by msg_id
- Always-on kernel status observation (IOPub)
- Execute requests with streaming output
- Input (stdin) handling
- Interrupt via message (for WebSocket-based interruption)
- Output size limit enforcement

Architecture:
    connect() starts a background reader task that continuously reads
    from the WebSocket. All IOPub status messages fire on_status_change
    regardless of whether an execution is active. Execution-related
    messages are routed to per-execution queues keyed by parent_msg_id.
    execute() creates a queue, sends the request, and reads from its
    queue until completion. Connection loss is propagated to all active
    executions via a disconnect sentinel.
"""

import asyncio
import json
import logging
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, AsyncIterator, Awaitable, Callable, Dict, Optional

from websockets.asyncio.client import connect as ws_connect, ClientConnection
from websockets.exceptions import ConnectionClosed

logger = logging.getLogger("skop-runner")

# Maximum accumulated output size per execution (10MB)
MAX_OUTPUT_SIZE = 10 * 1024 * 1024

# Default execution timeout (5 minutes)
DEFAULT_EXECUTION_TIMEOUT = 300.0

# WebSocket connection timeout
WS_CONNECT_TIMEOUT = 30.0

# Message receive timeout for polling
MSG_RECEIVE_TIMEOUT = 1.0

# Seconds of silence before checking kernel health via REST API
KERNEL_HEALTH_CHECK_INTERVAL = 10.0

# Sentinel value placed in execution queues when the WebSocket disconnects.
# Consumers check `msg is _DISCONNECT_SENTINEL` to detect connection loss.
_DISCONNECT_SENTINEL = object()


# =============================================================================
# Exceptions
# =============================================================================


class JupyterWebSocketError(Exception):
    """Base exception for Jupyter WebSocket errors."""

    pass


class JupyterWebSocketConnectionError(JupyterWebSocketError):
    """Failed to connect or connection lost."""

    pass


class JupyterKernelDiedError(JupyterWebSocketError):
    """Kernel process died during execution."""

    pass


class JupyterExecutionTimeoutError(JupyterWebSocketError):
    """Execution exceeded timeout."""

    pass


# =============================================================================
# Message Building
# =============================================================================


def _make_header(msg_type: str, session: str, username: str = "skop-runner") -> Dict[str, Any]:
    """Create a Jupyter message header."""
    return {
        "msg_id": str(uuid.uuid4()),
        "msg_type": msg_type,
        "username": username,
        "session": session,
        "date": datetime.now(timezone.utc).isoformat(),
        "version": "5.3",
    }


def _make_message(
    msg_type: str,
    content: Dict[str, Any],
    session: str,
    parent_header: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """Create a complete Jupyter protocol message."""
    return {
        "header": _make_header(msg_type, session),
        "parent_header": parent_header or {},
        "metadata": {},
        "content": content,
        "buffers": [],
        "channel": "shell",
    }


# =============================================================================
# Pending Execution Tracker
# =============================================================================


@dataclass
class PendingExecution:
    """Tracks state for a pending execute_request."""

    msg_id: str
    code: str
    start_time: float
    timeout: float
    total_output_size: int = 0
    output_truncated: bool = False
    status: str = "pending"  # pending, ok, error
    completed: bool = False
    awaiting_input: bool = False
    input_value: Optional[str] = None  # Set when input is provided
    input_parent_header: Optional[Dict[str, Any]] = None  # Header for input reply


# =============================================================================
# WebSocket Client
# =============================================================================


class JupyterKernelWebSocket:
    """
    WebSocket client for Jupyter kernel communication.

    Connects to a single kernel's WebSocket channel and provides methods
    for execution, input handling, and output streaming. A background
    reader task continuously reads from the WebSocket, routing messages
    to per-execution queues and always firing on_status_change for IOPub
    status messages.

    Example:
        ws = JupyterKernelWebSocket(kernel_id, "http://127.0.0.1:18888", "token")
        await ws.connect()   # starts background reader
        async for output in ws.execute("print('hello')"):
            print(output)
        await ws.close()      # stops background reader
    """

    def __init__(
        self,
        kernel_id: str,
        base_url: str,
        token: str,
        on_kernel_died: Optional[Callable[[], None]] = None,
        is_kernel_alive: Optional[Callable[[], Awaitable[bool]]] = None,
        on_status_change: Optional[Callable[[str], None]] = None,
    ):
        """
        Initialize the WebSocket client.

        Args:
            kernel_id: The kernel UUID.
            base_url: Jupyter Server base URL (e.g., "http://127.0.0.1:18888").
            token: Authentication token.
            on_kernel_died: Optional callback when kernel dies.
            is_kernel_alive: Optional async callable returning True if kernel is alive.
            on_status_change: Optional callback when kernel execution state changes
                (called with "busy", "idle", or "starting").
        """
        self._kernel_id = kernel_id
        self._base_url = base_url.rstrip("/")
        self._token = token
        self._on_kernel_died = on_kernel_died
        self._is_kernel_alive = is_kernel_alive
        self._on_status_change = on_status_change

        # WebSocket connection
        self._ws: Optional[ClientConnection] = None
        self._session = str(uuid.uuid4())

        # Message routing
        self._pending_executions: Dict[str, PendingExecution] = {}

        # Background reader state
        self._execution_queues: Dict[str, asyncio.Queue] = {}
        self._reader_task: Optional[asyncio.Task] = None
        self._closing: bool = False

        # Connection state
        self._connected = False

        # Lock for connection state
        self._lock = asyncio.Lock()

    @property
    def kernel_id(self) -> str:
        """Get the kernel ID."""
        return self._kernel_id

    @property
    def connected(self) -> bool:
        """Check if WebSocket is connected."""
        return self._connected and self._ws is not None

    def _build_ws_url(self) -> str:
        """Build the WebSocket URL for kernel channels."""
        ws_base = self._base_url.replace("http://", "ws://").replace(
            "https://", "wss://"
        )
        return f"{ws_base}/api/kernels/{self._kernel_id}/channels?token={self._token}"

    async def connect(self) -> None:
        """
        Connect to the kernel WebSocket.

        Raises:
            JupyterWebSocketConnectionError: If connection fails.
        """
        async with self._lock:
            if self._connected:
                return

            ws_url = self._build_ws_url()
            logger.debug(f"Connecting to kernel WebSocket: {self._kernel_id}")

            try:
                self._ws = await asyncio.wait_for(
                    ws_connect(ws_url),
                    timeout=WS_CONNECT_TIMEOUT,
                )
                self._connected = True
                self._closing = False
                self._reader_task = asyncio.create_task(
                    self._background_reader(),
                    name=f"ws-reader-{self._kernel_id}",
                )
                logger.info(f"Connected to kernel WebSocket: {self._kernel_id}")
            except asyncio.TimeoutError:
                raise JupyterWebSocketConnectionError(
                    f"Timeout connecting to kernel {self._kernel_id}"
                )
            except Exception as e:
                raise JupyterWebSocketConnectionError(
                    f"Failed to connect to kernel {self._kernel_id}: {e}"
                ) from e

    async def close(self) -> None:
        """Close the WebSocket connection."""
        async with self._lock:
            if self._ws is not None:
                self._closing = True
                try:
                    await self._ws.close()
                except Exception as e:
                    logger.debug(f"Error closing WebSocket: {e}")
                finally:
                    self._ws = None
                    self._connected = False

                    # Stop background reader
                    if self._reader_task and not self._reader_task.done():
                        self._reader_task.cancel()
                        try:
                            await asyncio.wait_for(
                                self._reader_task, timeout=5.0
                            )
                        except (asyncio.CancelledError, asyncio.TimeoutError):
                            pass
                    self._reader_task = None

                    # Clear execution queues
                    self._execution_queues.clear()

                    logger.info(f"Closed kernel WebSocket: {self._kernel_id}")

    async def __aenter__(self) -> "JupyterKernelWebSocket":
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        await self.close()

    async def _background_reader(self) -> None:
        """
        Continuously read from the WebSocket and route messages.

        This task runs for the lifetime of the connection. It:
        1. Always fires on_status_change for IOPub status messages
        2. Routes execution-related messages to per-execution queues
        3. Detects connection loss and notifies all consumers
        """
        try:
            while True:
                try:
                    raw = await self._ws.recv()
                except ConnectionClosed:
                    break
                except Exception as e:
                    logger.warning(f"Unexpected error reading WebSocket: {e}")
                    break

                # Parse message
                try:
                    msg = json.loads(raw)
                except json.JSONDecodeError as e:
                    logger.warning(f"Invalid JSON from kernel: {e}")
                    continue

                channel = msg.get("channel", "")
                msg_type = msg.get("header", {}).get("msg_type", "")
                content = msg.get("content", {})
                parent_msg_id = msg.get("parent_header", {}).get("msg_id")

                # Always observe IOPub status changes
                if channel == "iopub" and msg_type == "status":
                    execution_state = content.get("execution_state", "")

                    # Jupyter Server sends "dead" when the kernel process exits.
                    # This is a death signal, not a status transition — handle it
                    # exclusively via on_kernel_died (in the finally block), not
                    # via on_status_change.
                    if execution_state == "dead":
                        logger.warning(
                            f"Kernel {self._kernel_id} reported dead via iopub"
                        )
                        break

                    if self._on_status_change and execution_state:
                        try:
                            self._on_status_change(execution_state)
                        except Exception as e:
                            logger.warning(
                                f"on_status_change callback failed: {e}"
                            )

                # Route to execution queue if one is registered
                if parent_msg_id:
                    queue = self._execution_queues.get(parent_msg_id)
                    if queue is not None:
                        await queue.put(msg)

        finally:
            # Connection lost or closing
            self._connected = False

            # Notify all active executions via sentinel
            for queue in self._execution_queues.values():
                try:
                    await queue.put(_DISCONNECT_SENTINEL)
                except Exception:
                    pass

            # Fire on_kernel_died only for unexpected disconnects
            if not self._closing and self._on_kernel_died:
                try:
                    self._on_kernel_died()
                except Exception as e:
                    logger.warning(f"on_kernel_died callback failed: {e}")

            logger.debug(
                f"Background reader stopped for kernel {self._kernel_id} "
                f"(closing={self._closing})"
            )

    async def _send_message(self, msg: Dict[str, Any]) -> None:
        """Send a message over the WebSocket."""
        if not self._connected or self._ws is None:
            raise JupyterWebSocketConnectionError("WebSocket not connected")

        try:
            await self._ws.send(json.dumps(msg))
        except ConnectionClosed as e:
            self._connected = False
            raise JupyterWebSocketConnectionError(
                f"Connection closed while sending: {e}"
            ) from e

    def _make_execute_request(self, code: str) -> Dict[str, Any]:
        """Create an execute_request message."""
        content = {
            "code": code,
            "silent": False,
            "store_history": True,
            "user_expressions": {},
            "allow_stdin": True,
            "stop_on_error": True,
        }
        msg = _make_message("execute_request", content, self._session)
        msg["channel"] = "shell"
        return msg

    def _make_input_reply(self, value: str, parent_header: Dict[str, Any]) -> Dict[str, Any]:
        """Create an input_reply message."""
        content = {"value": value}
        msg = _make_message("input_reply", content, self._session, parent_header)
        msg["channel"] = "stdin"
        return msg

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

        Args:
            code: Code to execute.
            timeout: Maximum execution time in seconds.

        Yields:
            Output dicts matching kernel_manager.py format:
            - stream: {"output_type": "stream", "name": "stdout/stderr", "text": "..."}
            - execute_result: {"output_type": "execute_result", "data": {...}, ...}
            - display_data: {"output_type": "display_data", "data": {...}, ...}
            - update_display_data: {"output_type": "update_display_data", ...}
            - clear_output: {"output_type": "clear_output", "wait": bool}
            - error: {"output_type": "error", "ename": "...", ...}
            - input_request: {"output_type": "input_request", "prompt": "...", ...}

        Raises:
            JupyterWebSocketConnectionError: If connection is lost.
            JupyterExecutionTimeoutError: If execution times out.
        """
        if not self._connected:
            await self.connect()

        # Send execute request
        request = self._make_execute_request(code)
        msg_id = request["header"]["msg_id"]

        # Track pending execution
        pending = PendingExecution(
            msg_id=msg_id,
            code=code,
            start_time=asyncio.get_running_loop().time(),
            timeout=timeout,
        )
        self._pending_executions[msg_id] = pending

        # Create execution queue for background reader to route messages to
        execution_queue = asyncio.Queue()
        self._execution_queues[msg_id] = execution_queue

        try:
            await self._send_message(request)
            logger.debug(f"Sent execute_request: {msg_id}")

            # Process messages until execution completes
            async for output in self._process_execution(pending, execution_queue):
                yield output

        finally:
            # Clean up pending execution and queue
            self._pending_executions.pop(msg_id, None)
            self._execution_queues.pop(msg_id, None)

    async def _process_execution(
        self, pending: PendingExecution, queue: asyncio.Queue
    ) -> AsyncIterator[Dict[str, Any]]:
        """Process messages for a pending execution from its queue."""
        last_message_time = asyncio.get_running_loop().time()

        while not pending.completed:
            # Check timeout (0 = no timeout)
            elapsed = asyncio.get_running_loop().time() - pending.start_time
            if pending.timeout > 0 and elapsed > pending.timeout:
                logger.warning(
                    f"Execution timed out after {pending.timeout}s: {pending.msg_id}"
                )
                pending.completed = True
                pending.status = "error"
                yield {
                    "output_type": "error",
                    "ename": "TimeoutError",
                    "evalue": f"Execution timed out after {pending.timeout} seconds",
                    "traceback": [
                        f"TimeoutError: Cell execution exceeded {pending.timeout} second limit.",
                        "Consider breaking your code into smaller cells or optimizing performance.",
                    ],
                }
                return

            # Check if input was provided while we were reading messages
            if pending.awaiting_input and pending.input_value is not None:
                reply = self._make_input_reply(
                    pending.input_value, pending.input_parent_header
                )
                await self._send_message(reply)
                logger.debug(f"Sent input_reply for {pending.msg_id}")
                pending.awaiting_input = False
                pending.input_value = None
                pending.input_parent_header = None
                last_message_time = asyncio.get_running_loop().time()

            # Read next message from execution queue
            try:
                msg = await asyncio.wait_for(
                    queue.get(), timeout=MSG_RECEIVE_TIMEOUT
                )
            except asyncio.TimeoutError:
                msg = None

            # Check for disconnect sentinel
            if msg is _DISCONNECT_SENTINEL:
                pending.completed = True
                pending.status = "error"
                yield {
                    "output_type": "error",
                    "ename": "KernelDied",
                    "evalue": "Kernel connection lost",
                    "traceback": [
                        "KernelDied: The connection to the kernel was lost.",
                        "This may be due to the kernel crashing or network issues.",
                        "Please restart the kernel and try again.",
                    ],
                }
                # on_kernel_died already fired by background reader
                return

            if msg is None:
                # Check kernel health periodically during silence
                if (self._is_kernel_alive is not None
                        and not pending.awaiting_input):
                    now = asyncio.get_running_loop().time()
                    silence = now - last_message_time
                    if silence >= KERNEL_HEALTH_CHECK_INTERVAL:
                        try:
                            alive = await self._is_kernel_alive()
                        except Exception:
                            alive = True  # Assume alive on API error
                        if not alive:
                            pending.completed = True
                            pending.status = "error"
                            yield {
                                "output_type": "error",
                                "ename": "KernelDied",
                                "evalue": "Kernel process died",
                                "traceback": [
                                    "KernelDied: The kernel process has exited unexpectedly.",
                                    "Please restart the kernel and try again.",
                                ],
                            }
                            # Background reader may not have detected death yet
                            # (WS to Jupyter Server can stay open after kernel dies)
                            if self._on_kernel_died:
                                try:
                                    self._on_kernel_died()
                                except Exception:
                                    pass
                            return
                        # Kernel alive — reset timer
                        last_message_time = now
                continue

            # Reset timer on message receipt
            last_message_time = asyncio.get_running_loop().time()

            channel = msg.get("channel", "")
            msg_type = msg.get("header", {}).get("msg_type", "")
            content = msg.get("content", {})

            # Handle by channel and type
            if channel == "iopub":
                async for output in self._handle_iopub_message(
                    msg_type, content, msg, pending
                ):
                    yield output
            elif channel == "stdin":
                async for output in self._handle_stdin_message(
                    msg_type, content, msg, pending
                ):
                    yield output
            elif channel == "shell":
                # Shell channel contains execute_reply (completion signal)
                if msg_type == "execute_reply":
                    # Don't mark complete yet - wait for idle status
                    pending.status = content.get("status", "ok")

    async def _handle_iopub_message(
        self,
        msg_type: str,
        content: Dict[str, Any],
        msg: Dict[str, Any],
        pending: PendingExecution,
    ) -> AsyncIterator[Dict[str, Any]]:
        """Handle IOPub channel messages."""
        if msg_type == "stream":
            text = content.get("text", "")
            pending.total_output_size += len(text)

            if pending.total_output_size > MAX_OUTPUT_SIZE and not pending.output_truncated:
                pending.output_truncated = True
                yield {
                    "output_type": "stream",
                    "name": "stderr",
                    "text": f"\n[Output truncated: exceeded {MAX_OUTPUT_SIZE // (1024*1024)}MB limit]\n",
                }
            elif not pending.output_truncated:
                yield {
                    "output_type": "stream",
                    "name": content.get("name", "stdout"),
                    "text": text,
                }

        elif msg_type == "execute_result":
            data = content.get("data", {})
            self._track_data_size(data, pending)

            if pending.total_output_size > MAX_OUTPUT_SIZE and not pending.output_truncated:
                pending.output_truncated = True
                yield {
                    "output_type": "stream",
                    "name": "stderr",
                    "text": f"\n[Output truncated: exceeded {MAX_OUTPUT_SIZE // (1024*1024)}MB limit]\n",
                }
            elif not pending.output_truncated:
                yield {
                    "output_type": "execute_result",
                    "data": data,
                    "execution_count": content.get("execution_count", 0),
                    "metadata": content.get("metadata", {}),
                }

        elif msg_type == "display_data":
            data = content.get("data", {})
            self._track_data_size(data, pending)

            if pending.total_output_size > MAX_OUTPUT_SIZE and not pending.output_truncated:
                pending.output_truncated = True
                yield {
                    "output_type": "stream",
                    "name": "stderr",
                    "text": f"\n[Output truncated: exceeded {MAX_OUTPUT_SIZE // (1024*1024)}MB limit]\n",
                }
            elif not pending.output_truncated:
                display_id = content.get("transient", {}).get("display_id")
                result: Dict[str, Any] = {
                    "output_type": "display_data",
                    "data": data,
                    "metadata": content.get("metadata", {}),
                }
                if display_id:
                    result["display_id"] = display_id
                yield result

        elif msg_type == "update_display_data":
            display_id = content.get("transient", {}).get("display_id")
            if display_id:
                data = content.get("data", {})
                self._track_data_size(data, pending)

                if pending.total_output_size > MAX_OUTPUT_SIZE and not pending.output_truncated:
                    pending.output_truncated = True
                    yield {
                        "output_type": "stream",
                        "name": "stderr",
                        "text": f"\n[Output truncated: exceeded {MAX_OUTPUT_SIZE // (1024*1024)}MB limit]\n",
                    }
                elif not pending.output_truncated:
                    yield {
                        "output_type": "update_display_data",
                        "data": data,
                        "metadata": content.get("metadata", {}),
                        "display_id": display_id,
                    }

        elif msg_type == "clear_output":
            yield {
                "output_type": "clear_output",
                "wait": bool(content.get("wait", False)),
            }

        elif msg_type == "error":
            pending.status = "error"
            yield {
                "output_type": "error",
                "ename": content.get("ename", "Error"),
                "evalue": content.get("evalue", ""),
                "traceback": content.get("traceback", []),
            }

        elif msg_type == "status":
            execution_state = content.get("execution_state", "")
            # on_status_change is fired by background reader (not here)
            # to avoid double-firing during execution
            if execution_state == "idle":
                # Execution complete
                pending.completed = True

    async def _handle_stdin_message(
        self,
        msg_type: str,
        content: Dict[str, Any],
        msg: Dict[str, Any],
        pending: PendingExecution,
    ) -> AsyncIterator[Dict[str, Any]]:
        """Handle stdin channel messages (input requests).

        Non-blocking: yields the input_request and returns immediately.
        The actual input reply is sent from the _process_execution loop
        when send_input() sets pending.input_value.
        """
        if msg_type == "input_request":
            pending.awaiting_input = True
            pending.input_value = None
            pending.input_parent_header = msg.get("header", {})

            yield {
                "output_type": "input_request",
                "prompt": content.get("prompt", ""),
                "password": content.get("password", False),
            }

    def _track_data_size(self, data: Dict[str, Any], pending: PendingExecution) -> None:
        """Track output size from data dict."""
        for value in data.values():
            if isinstance(value, str):
                pending.total_output_size += len(value)

    async def send_input(self, value: str) -> None:
        """
        Send input reply for a pending input request.

        Args:
            value: The input value to send.

        Raises:
            ValueError: If no execution is awaiting input.
        """
        # Find the pending execution that's awaiting input
        for pending in self._pending_executions.values():
            if pending.awaiting_input:
                pending.input_value = value
                return

        raise ValueError("No execution is awaiting input")

    @property
    def awaiting_input(self) -> bool:
        """Check if any execution is awaiting input."""
        return any(p.awaiting_input for p in self._pending_executions.values())
