"""Round-robin admission by user, with bounded pending work."""

from __future__ import annotations

import asyncio
import logging
import threading
import time
from collections import Counter, deque
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from dataclasses import dataclass

logger = logging.getLogger(__name__)


class QueueFullError(RuntimeError):
    pass


class FairAsyncGate:
    def __init__(self, capacity: int, per_user: int, queue_limit: int = 256):
        self.capacity, self.per_user, self.queue_limit = capacity, min(capacity, per_user), queue_limit
        self._active = Counter()
        self._queues: dict[str, deque] = {}
        self._order = deque()

    def _dispatch(self):
        while sum(self._active.values()) < self.capacity:
            for _ in range(len(self._order)):
                key = self._order.popleft()
                if self._active[key] < self.per_user:
                    future = self._queues[key].popleft()
                    if self._queues[key]:
                        self._order.append(key)
                    else:
                        del self._queues[key]
                    self._active[key] += 1
                    future.set_result(None)
                    break
                self._order.append(key)
            else:
                return

    def _release(self, key):
        self._active[key] -= 1
        if not self._active[key]:
            del self._active[key]
        self._dispatch()

    @asynccontextmanager
    async def slot(self, key: str):
        if len(self._queues.get(key, ())) >= max(1, self.queue_limit // max(2, self.capacity)) or sum(len(queue) for queue in self._queues.values()) >= self.queue_limit:
            raise QueueFullError("Background queue is full; retry later")
        future = asyncio.get_running_loop().create_future()
        if key not in self._queues:
            self._queues[key] = deque()
            self._order.append(key)
        self._queues[key].append(future)
        self._dispatch()
        try:
            await asyncio.shield(future)
        except BaseException:
            if future.done():
                self._release(key)
            else:
                self._queues[key].remove(future)
                if not self._queues[key]:
                    del self._queues[key]
                    self._order.remove(key)
                future.cancel()
                self._dispatch()
            raise
        try:
            yield
        finally:
            self._release(key)


@dataclass
class _Job:
    key: str
    run: Callable[[float], None]
    deadline: float
    cancelled: Callable[[], bool]
    expired: Callable[[bool], None]


class FairBackgroundScheduler:
    """Only admitted work enters the executor, so queued jobs occupy no worker.

    Deadlines include queue time. Expiry/cancellation never releases the slot of
    a running job; its execution wrapper must first stop the actual work.
    """

    def __init__(self, capacity=3, per_user=2, queue_limit=256):
        self.capacity, self.per_user, self.queue_limit = capacity, min(capacity, per_user), queue_limit
        self._condition = threading.Condition()
        self._queues: dict[str, deque[_Job]] = {}
        self._order = deque()
        self._active = Counter()
        self._closed = False
        self._pool = ThreadPoolExecutor(max_workers=capacity, thread_name_prefix="subagent-fair")
        self._dispatcher = threading.Thread(target=self._dispatch, name="subagent-admission", daemon=True)
        self._dispatcher.start()

    def submit(self, key, run, *, timeout, cancelled, expired):
        with self._condition:
            if self._closed or len(self._queues.get(key, ())) >= max(1, self.queue_limit // max(2, self.capacity)) or sum(map(len, self._queues.values())) >= self.queue_limit:
                raise QueueFullError("Subagent queue is full; retry later")
            if key not in self._queues:
                self._queues[key] = deque()
                self._order.append(key)
            self._queues[key].append(_Job(key, run, time.monotonic() + timeout, cancelled, expired))
            self._condition.notify_all()

    def _execute(self, job):
        try:
            remaining = job.deadline - time.monotonic()
            if job.cancelled() or remaining <= 0:
                self._expire(job, job.cancelled())
            else:
                job.run(remaining)
        finally:
            with self._condition:
                self._active[job.key] -= 1
                if not self._active[job.key]:
                    del self._active[job.key]
                self._condition.notify_all()

    @staticmethod
    def _expire(job, cancelled):
        try:
            job.expired(cancelled)
        except Exception:
            logger.exception("Background expiry callback failed for user %s", job.key)

    def _dispatch(self):
        with self._condition:
            while not self._closed:
                for key in list(self._order):
                    retained = deque()
                    for job in self._queues[key]:
                        if job.cancelled() or job.deadline <= time.monotonic():
                            self._expire(job, job.cancelled())
                        else:
                            retained.append(job)
                    if retained:
                        self._queues[key] = retained
                    else:
                        del self._queues[key]
                        self._order.remove(key)
                admitted = False
                if sum(self._active.values()) < self.capacity:
                    for _ in range(len(self._order)):
                        key = self._order.popleft()
                        if self._active[key] < self.per_user:
                            job = self._queues[key].popleft()
                            if self._queues[key]:
                                self._order.append(key)
                            else:
                                del self._queues[key]
                            self._active[key] += 1
                            self._pool.submit(self._execute, job)
                            admitted = True
                            break
                        self._order.append(key)
                if not admitted:
                    self._condition.wait(timeout=0.05)

    def shutdown(self, wait=True):
        with self._condition:
            self._closed = True
            for queue in self._queues.values():
                for job in queue:
                    self._expire(job, True)
            self._queues.clear()
            self._order.clear()
            self._condition.notify_all()
        self._dispatcher.join()
        self._pool.shutdown(wait=wait)