"""
Async queue with concurrency control, deduplication, and observer callbacks.

Python equivalent of ts-agent's Queue<T> (src/libs/queue/queue.ts).
"""

from __future__ import annotations

import asyncio
import logging
from collections.abc import Awaitable, Callable
from typing import Generic, TypeVar

logger = logging.getLogger(__name__)

T = TypeVar("T")


class Queue(Generic[T]):
    def __init__(self, concurrency: int, get_item_key: Callable[[T], str]) -> None:
        self._semaphore = asyncio.Semaphore(concurrency)
        self._concurrency = concurrency
        self._get_item_key = get_item_key
        self._active_keys: list[str] = []
        self._pending_tasks: set[asyncio.Task[None]] = set()
        self._abort_events: dict[str, asyncio.Event] = {}
        self._on_start_listeners: list[Callable[[str], None]] = []
        self._on_done_listeners: list[Callable[[str], None]] = []
        self._on_added_listeners: list[Callable[[str], None]] = []
        self._on_aborted_listeners: list[Callable[[str], None]] = []

    # ── Enqueueing ──────────────────────────────────────────────────────────

    def enqueue(self, process: T, task: Callable[[], Awaitable[None]]) -> asyncio.Task[None] | None:
        """
        Synchronously submits a task to the queue. Returns the asyncio.Task,
        or None if the key is already active (deduplication).

        Callers that want to await completion should do:
            t = queue.enqueue(process, task)
            if t: await t
        """
        key = self._get_item_key(process)
        if key in self._active_keys:
            return None

        abort_event = asyncio.Event()
        self._abort_events[key] = abort_event
        self._active_keys.append(key)
        for fn in self._on_added_listeners:
            fn(key)

        async def _run() -> None:
            try:
                async with self._semaphore:
                    for fn in self._on_start_listeners:
                        fn(key)
                    if not abort_event.is_set():
                        await task()
            except asyncio.CancelledError:
                pass
            except Exception:
                logger.exception("Queue task failed for key: %s", key)
            finally:
                if key in self._active_keys:
                    self._active_keys.remove(key)
                self._abort_events.pop(key, None)
                for fn in self._on_done_listeners:
                    fn(key)

        t: asyncio.Task[None] = asyncio.create_task(_run())
        self._pending_tasks.add(t)
        t.add_done_callback(self._pending_tasks.discard)
        return t

    async def add(self, process: T, task: Callable[[], Awaitable[None]]) -> None:
        """Enqueue a task and await its completion."""
        t = self.enqueue(process, task)
        if t is not None:
            await t

    # ── State inspection ────────────────────────────────────────────────────

    def has(self, process: T) -> bool:
        return self._get_item_key(process) in self._active_keys

    def get_active_processes(self) -> list[str]:
        return list(self._active_keys)

    def get_processing(self) -> list[str]:
        """Tasks currently running (up to concurrency limit)."""
        return self._active_keys[: self._concurrency]

    def get_queued(self) -> list[str]:
        """Tasks waiting for a concurrency slot."""
        return self._active_keys[self._concurrency :]

    def should_queue(self) -> bool:
        return len(self._active_keys) >= self._concurrency

    # ── Observer registration ───────────────────────────────────────────────

    def on_start(self, fn: Callable[[str], None]) -> None:
        """Fired every time a task starts executing (acquires semaphore)."""
        self._on_start_listeners.append(fn)

    def on_done(self, fn: Callable[[str], None]) -> None:
        """Fired every time a task finishes (success or error)."""
        self._on_done_listeners.append(fn)

    def on_added(self, fn: Callable[[str], None]) -> None:
        """Fired every time a task is enqueued."""
        self._on_added_listeners.append(fn)

    def on_aborted(self, fn: Callable[[str], None]) -> None:
        """Fired every time a task is aborted."""
        self._on_aborted_listeners.append(fn)

    # ── Idle / wait ─────────────────────────────────────────────────────────

    async def get_idle_promise(self) -> None:
        """Resolves when all currently pending tasks have completed."""
        while self._pending_tasks:
            await asyncio.gather(*list(self._pending_tasks), return_exceptions=True)

    # ── Abort ───────────────────────────────────────────────────────────────

    def abort(self, process: T) -> None:
        key = self._get_item_key(process)
        if key in self._active_keys:
            self._active_keys.remove(key)
        event = self._abort_events.pop(key, None)
        if event:
            event.set()
        for fn in self._on_aborted_listeners:
            fn(key)

    def abort_all(self) -> None:
        keys = list(self._active_keys)
        self._active_keys.clear()
        for key in keys:
            event = self._abort_events.pop(key, None)
            if event:
                event.set()
            for fn in self._on_aborted_listeners:
                fn(key)
        self._abort_events.clear()
        for t in list(self._pending_tasks):
            t.cancel()
        self._pending_tasks.clear()
