from __future__ import annotations

import asyncio
import heapq
import logging
from collections.abc import Awaitable, Callable

logger = logging.getLogger(__name__)


class OrderedEventQueue:
    def __init__(self, max_size: int = 1000, max_wait_ms: int = 5000):
        self._max_size = max_size
        self._max_wait_ms = max_wait_ms
        self._heap: list[tuple[int, int, dict]] = []
        self._seq_counter = 0
        self._locked = False
        self._consumer_task: asyncio.Task | None = None
        self._handler: Callable[[dict], Awaitable[None]] | None = None
        self._event = asyncio.Event()

    def set_handler(self, handler: Callable[[dict], Awaitable[None]]) -> None:
        self._handler = handler

    async def push(self, event: dict, timestamp: int = 0) -> None:
        self._seq_counter += 1
        entry = (timestamp, self._seq_counter, event)
        heapq.heappush(self._heap, entry)
        if len(self._heap) > self._max_size:
            heapq.heappop(self._heap)
        self._event.set()

    async def start(self) -> None:
        if self._consumer_task is not None:
            return
        self._consumer_task = asyncio.create_task(self._consume_loop())

    async def stop(self) -> None:
        if self._consumer_task:
            self._consumer_task.cancel()
            try:
                await self._consumer_task
            except asyncio.CancelledError:
                pass
            self._consumer_task = None

    async def _consume_loop(self) -> None:
        while True:
            try:
                if not self._heap:
                    await asyncio.wait_for(
                        self._event.wait(),
                        timeout=self._max_wait_ms / 1000.0,
                    )
                    self._event.clear()
                    if not self._heap:
                        continue

                ts, seq, event = heapq.heappop(self._heap)
                if self._handler:
                    try:
                        await self._handler(event)
                    except Exception:
                        logger.exception("OrderedEventQueue handler failed")

            except TimeoutError:
                continue
            except asyncio.CancelledError:
                break
            except Exception:
                logger.exception("OrderedEventQueue consume loop error")

    @property
    def size(self) -> int:
        return len(self._heap)