"""In-memory run queue for gateway-driven agent turns."""

from __future__ import annotations

from collections import deque
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass
import threading
from typing import Callable, Generic, TypeVar


T = TypeVar("T")


@dataclass(slots=True)
class _QueueEntry(Generic[T]):
    future: Future[T]
    task: Callable[[], T]


class InMemoryRunQueue:
    """Serialize agent turns per session with a global concurrency cap."""

    def __init__(self, *, max_concurrent: int = 1) -> None:
        self.max_concurrent = max(1, int(max_concurrent))
        self._lock = threading.RLock()
        self._executor = ThreadPoolExecutor(
            max_workers=self.max_concurrent,
            thread_name_prefix="memnixa-run-queue",
        )
        self._queues: dict[str, deque[_QueueEntry[object]]] = {}
        self._ready_sessions: deque[str] = deque()
        self._ready_session_set: set[str] = set()
        self._active_sessions: set[str] = set()
        self._active_count = 0
        self._shutdown = False

    def enqueue(self, session_id: str, task: Callable[[], T]) -> Future[T]:
        """Queue one synchronous task and return a future for its result."""
        future: Future[T] = Future()
        cleaned_session_id = (session_id or "main").strip() or "main"
        with self._lock:
            if self._shutdown:
                future.set_exception(RuntimeError("run queue is shut down"))
                return future
            queue = self._queues.setdefault(cleaned_session_id, deque())
            queue.append(_QueueEntry(future=future, task=task))
            if (
                cleaned_session_id not in self._active_sessions
                and cleaned_session_id not in self._ready_session_set
            ):
                self._ready_sessions.append(cleaned_session_id)
                self._ready_session_set.add(cleaned_session_id)
            self._dispatch_locked()
        return future

    def snapshot(self) -> dict[str, object]:
        """Return queue activity counters for diagnostics."""
        with self._lock:
            return {
                "active_count": self._active_count,
                "active_sessions": sorted(self._active_sessions),
                "queued_count": sum(len(items) for items in self._queues.values()),
                "queued_by_session": {
                    session_id: len(items)
                    for session_id, items in sorted(self._queues.items())
                    if items
                },
                "max_concurrent": self.max_concurrent,
            }

    def shutdown(self, *, wait: bool = True) -> None:
        """Stop accepting new jobs and close worker threads."""
        with self._lock:
            self._shutdown = True
        self._executor.shutdown(wait=wait, cancel_futures=False)

    def _dispatch_locked(self) -> None:
        while self._active_count < self.max_concurrent and self._ready_sessions:
            session_id = self._ready_sessions.popleft()
            self._ready_session_set.discard(session_id)
            queue = self._queues.get(session_id)
            if not queue:
                self._queues.pop(session_id, None)
                continue
            entry = queue.popleft()
            if not queue:
                self._queues.pop(session_id, None)
            else:
                self._queues[session_id] = queue
            self._active_sessions.add(session_id)
            self._active_count += 1
            self._executor.submit(self._run_entry, session_id, entry)

    def _run_entry(self, session_id: str, entry: _QueueEntry[object]) -> None:
        try:
            result = entry.task()
        except Exception as exc:
            if not entry.future.done():
                entry.future.set_exception(exc)
        else:
            if not entry.future.done():
                entry.future.set_result(result)
        finally:
            with self._lock:
                self._active_sessions.discard(session_id)
                self._active_count = max(0, self._active_count - 1)
                if session_id in self._queues and session_id not in self._ready_session_set:
                    self._ready_sessions.append(session_id)
                    self._ready_session_set.add(session_id)
                self._dispatch_locked()
