from __future__ import annotations
import asyncio
import contextlib
from dataclasses import dataclass
from typing import Awaitable, Callable, Optional, Tuple

from .events import BusEvent

Handler = Callable[[BusEvent], Awaitable[None]]

@dataclass(order=True)
class _PQItem:
    priority: int
    seq: int
    event: BusEvent

class PriorityScheduler:
    """Async priority scheduler with backpressure and graceful shutdown."""

    def __init__(self, maxsize: int = 10_000):
        self._pq: asyncio.PriorityQueue[_PQItem] = asyncio.PriorityQueue(maxsize=maxsize)
        self._running = False
        self._task: Optional[asyncio.Task] = None
        self._seq = 0
        self._handler: Optional[Handler] = None

    def set_handler(self, handler: Handler) -> None:
        self._handler = handler

    async def put(self, event: BusEvent) -> None:
        # Lower number => higher priority
        self._seq += 1
        await self._pq.put(_PQItem(priority=int(event.priority), seq=self._seq, event=event))

    async def start(self) -> None:
        if self._running:
            return
        if self._handler is None:
            raise RuntimeError("PriorityScheduler handler not set.")
        self._running = True
        self._task = asyncio.create_task(self._loop())

    async def stop(self) -> None:
        self._running = False
        if self._task:
            self._task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._task

    async def _loop(self) -> None:
        assert self._handler is not None
        while self._running:
            item = await self._pq.get()
            await self._handler(item.event)
