"""Execution queue for serializing cell execution per kernel."""

from __future__ import annotations

import asyncio
import time
import logging
from dataclasses import dataclass, field
from typing import Callable, Awaitable, Optional

from .structured_log import get_logger as _get_slog

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


class ExecutionCancelledError(Exception):
    """Raised when an execution is cancelled by the user or its kernel dies.

    This is deliberately a regular Exception (not asyncio.CancelledError) so that
    it is caught by the generic ``except Exception`` handlers in handle_message
    and _handle_and_respond.  asyncio.CancelledError became a BaseException in
    Python 3.9, which would escape those handlers and silently drop the JSON-RPC
    response, leaving the frontend hanging forever.
    """

    pass


@dataclass
class QueuedExecution:
    """A single execution waiting in or running from the queue."""

    cell_ids: list[str]
    kernel_id: str
    execute_fn: Callable[[], Awaitable[dict]]
    future: asyncio.Future
    queued_at: float = field(default_factory=time.time)


class ExecutionQueue:
    """
    Per-kernel execution queue.

    Ensures only one execution runs at a time per kernel while allowing
    different kernels to execute in parallel.  The queue is a simple
    serializer — it doesn't know about cell details, just ensures one
    execution at a time.
    """

    def __init__(self) -> None:
        self._queues: dict[str, asyncio.Queue[QueuedExecution]] = {}
        self._workers: dict[str, asyncio.Task] = {}
        self._current: dict[str, QueuedExecution] = {}

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def submit(
        self,
        kernel_id: str,
        cell_ids: list[str],
        execute_fn: Callable[[], Awaitable[dict]],
    ) -> dict:
        """
        Submit an execution to the queue and wait for the result.

        Lazy-creates the queue + worker for the kernel if needed.
        Returns the result of ``execute_fn()``.
        Raises ``ExecutionCancelledError`` if cancelled before running.
        """
        loop = asyncio.get_running_loop()
        future: asyncio.Future[dict] = loop.create_future()

        item = QueuedExecution(
            cell_ids=cell_ids,
            kernel_id=kernel_id,
            execute_fn=execute_fn,
            future=future,
        )

        # Lazy-create queue + worker
        if kernel_id not in self._queues:
            self._queues[kernel_id] = asyncio.Queue()
        if kernel_id not in self._workers or self._workers[kernel_id].done():
            self._workers[kernel_id] = asyncio.create_task(
                self._worker(kernel_id),
                name=f"exec-worker-{kernel_id}",
            )

        await self._queues[kernel_id].put(item)
        return await future

    async def cancel_queue(
        self,
        kernel_id: str,
        interrupt_fn: Optional[Callable[[str], Awaitable[None]]] = None,
    ) -> dict:
        """
        Cancel all pending executions for a kernel and optionally interrupt
        the currently running one.

        Returns ``{"cancelled_cells": [...], "interrupted_cells": [...]}``.
        """
        cancelled_cells: list[str] = []
        interrupted_cells: list[str] = []

        # Drain pending items
        queue = self._queues.get(kernel_id)
        if queue:
            while not queue.empty():
                try:
                    item = queue.get_nowait()
                    cancelled_cells.extend(item.cell_ids)
                    if not item.future.done():
                        item.future.set_exception(
                            ExecutionCancelledError("Execution was cancelled")
                        )
                except asyncio.QueueEmpty:
                    break

        # Interrupt the currently running item
        current = self._current.get(kernel_id)
        if current and interrupt_fn:
            interrupted_cells.extend(current.cell_ids)
            try:
                await interrupt_fn(kernel_id)
            except Exception as exc:
                logger.warning(f"Interrupt failed for kernel {kernel_id}: {exc}")

        return {
            "cancelled_cells": cancelled_cells,
            "interrupted_cells": interrupted_cells,
        }

    def get_queue_state(self, kernel_id: str | None = None) -> dict:
        """
        Return queue state for one or all kernels.

        Returns ``{kernel_id: {"running": [...], "queued": N}}``.
        """
        result: dict[str, dict] = {}

        kernel_ids = [kernel_id] if kernel_id else list(self._queues.keys())
        for kid in kernel_ids:
            current = self._current.get(kid)
            queue = self._queues.get(kid)
            result[kid] = {
                "running": current.cell_ids if current else [],
                "queued": queue.qsize() if queue else 0,
            }

        return result

    async def cleanup_kernel(self, kernel_id: str) -> None:
        """
        Cancel the worker, drain the queue, and reject all futures
        for a kernel.  Called when a kernel dies.
        """
        rejected_count = 0

        # Reject the currently running item's future BEFORE cancelling
        # the worker — this unblocks any caller awaiting submit().
        current = self._current.pop(kernel_id, None)
        if current and not current.future.done():
            current.future.set_exception(
                ExecutionCancelledError("Kernel died")
            )
            rejected_count += 1

        # Cancel worker
        worker = self._workers.pop(kernel_id, None)
        if worker and not worker.done():
            worker.cancel()
            try:
                await worker
            except (asyncio.CancelledError, Exception):
                pass

        # Drain queue and reject futures
        queue = self._queues.pop(kernel_id, None)
        if queue:
            while not queue.empty():
                try:
                    item = queue.get_nowait()
                    if not item.future.done():
                        item.future.set_exception(
                            ExecutionCancelledError("Kernel died")
                        )
                        rejected_count += 1
                except asyncio.QueueEmpty:
                    break

        _log.info("queue.cleanup", kernel_id=kernel_id, rejected_count=rejected_count)

    # ------------------------------------------------------------------
    # Worker loop
    # ------------------------------------------------------------------

    async def _worker(self, kernel_id: str) -> None:
        """Infinite loop: pull items from the queue and execute them."""
        queue = self._queues[kernel_id]

        try:
            while True:
                item = await queue.get()

                # If the future was already cancelled/rejected (e.g. cleanup),
                # skip execution.
                if item.future.done():
                    continue

                self._current[kernel_id] = item
                try:
                    result = await item.execute_fn()
                    if not item.future.done():
                        item.future.set_result(result)
                except asyncio.CancelledError:
                    # CancelledError from worker.cancel() — reject future
                    # with ExecutionCancelledError (a regular Exception) so
                    # handle_message can catch it, then re-raise the real
                    # CancelledError to exit the worker loop.
                    if not item.future.done():
                        item.future.set_exception(
                            ExecutionCancelledError("Execution was cancelled")
                        )
                    raise
                except Exception as exc:
                    if not item.future.done():
                        item.future.set_exception(exc)
                finally:
                    self._current.pop(kernel_id, None)
        except asyncio.CancelledError:
            # Worker was cancelled (by cleanup_kernel) — exit cleanly
            pass
