"""FIFO job queue with concurrency semaphore.

Migrated pattern from old executor — FIFO queue + asyncio Semaphore
to limit concurrent executions (scripts & tools).

Architecture reference: Section 7.3 — Concurrency Model.
"""

from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass
from typing import Any, Callable, Coroutine, Set

logger = logging.getLogger("qremotex.execution")


@dataclass
class Job:
    """A unit of work to be executed."""

    request_id: str
    job_type: str  # "script" | "tool"
    payload: dict[str, Any]
    timeout_sec: int


class JobQueue:
    """FIFO job queue with bounded concurrency.

    Jobs are enqueued and dispatched to a worker callback when a semaphore
    slot is available.  Supports cancellation of queued and running jobs.
    """

    def __init__(
        self,
        concurrency: int,
        worker: Callable[[Job], Coroutine[Any, Any, None]],
    ) -> None:
        self._concurrency = concurrency
        self._sem = asyncio.Semaphore(concurrency)
        self._worker = worker

        self._queue: asyncio.Queue[Job] = asyncio.Queue()
        self._queued_ids: set[str] = set()
        self._running_ids: set[str] = set()
        self._tasks: Set[asyncio.Task] = set()  # type: ignore[type-arg]
        self._dispatcher_task: asyncio.Task | None = None

    # ------------------------------------------------------------------ #
    # Lifecycle
    # ------------------------------------------------------------------ #
    def start(self) -> None:
        """Start the dispatcher background task."""
        self._dispatcher_task = asyncio.create_task(self._dispatcher())

    async def stop(self) -> None:
        """Stop the dispatcher and cancel all running tasks."""
        if self._dispatcher_task:
            self._dispatcher_task.cancel()
        for t in list(self._tasks):
            t.cancel()

    # ------------------------------------------------------------------ #
    # Enqueue / Cancel
    # ------------------------------------------------------------------ #
    async def enqueue(self, job: Job) -> int:
        """Add a job to the queue.  Returns queue position (1-based)."""
        await self._queue.put(job)
        self._queued_ids.add(job.request_id)
        return self._queue.qsize()

    async def cancel(self, request_id: str) -> bool:
        """Cancel a queued or running job.  Returns True if found."""
        # Try to remove from queue
        if request_id in self._queued_ids:
            drained: list[Job] = []
            removed = False
            while not self._queue.empty():
                job = await self._queue.get()
                if job.request_id == request_id and not removed:
                    removed = True
                    self._queued_ids.discard(request_id)
                    continue
                drained.append(job)
            for j in drained:
                await self._queue.put(j)
            if removed:
                return True

        # Try to cancel running task (task itself handles process cleanup)
        if request_id in self._running_ids:
            for t in list(self._tasks):
                if getattr(t, "_job_request_id", None) == request_id:
                    t.cancel()
                    return True
        return False

    # ------------------------------------------------------------------ #
    # Stats
    # ------------------------------------------------------------------ #
    @property
    def queued_count(self) -> int:
        return self._queue.qsize()

    @property
    def running_count(self) -> int:
        return len(self._running_ids)

    # ------------------------------------------------------------------ #
    # Internal
    # ------------------------------------------------------------------ #
    async def _dispatcher(self) -> None:
        """Pull jobs FIFO and launch them when a semaphore slot opens."""
        while True:
            job = await self._queue.get()
            self._queued_ids.discard(job.request_id)

            task = asyncio.create_task(self._run_with_semaphore(job))
            task._job_request_id = job.request_id  # type: ignore[attr-defined]
            self._tasks.add(task)
            task.add_done_callback(self._task_done)

    async def _run_with_semaphore(self, job: Job) -> None:
        """Acquire a semaphore slot, then run the worker."""
        async with self._sem:
            self._running_ids.add(job.request_id)
            try:
                await self._worker(job)
            except asyncio.CancelledError:
                logger.info("Job %s cancelled", job.request_id)
            except Exception:
                logger.exception("Job %s failed with unhandled error", job.request_id)
            finally:
                self._running_ids.discard(job.request_id)

    def _task_done(self, task: asyncio.Task) -> None:
        self._tasks.discard(task)
