"""Bounded per-edge queues with explicit overload behaviour."""

from __future__ import annotations

from collections import deque
from collections.abc import AsyncIterator
from dataclasses import dataclass
from enum import Enum
from typing import Deque, Generic, TypeVar

import asyncio


ItemT = TypeVar("ItemT")


class OverflowPolicy(str, Enum):
    """Action taken when a producer reaches an edge's capacity."""

    BLOCK = "block"
    DROP_OLDEST = "drop_oldest"
    DROP_NEWEST = "drop_newest"
    REJECT = "reject"


class PutResult(str, Enum):
    """Non-exceptional outcome of :meth:`EdgeQueue.put`."""

    ENQUEUED = "enqueued"
    DROPPED_OLDEST = "dropped_oldest"
    DROPPED_NEWEST = "dropped_newest"


class EdgeQueueError(RuntimeError):
    """Base class for edge queue failures."""


class EdgeQueueClosed(EdgeQueueError):
    """Raised when reading a drained closed queue or writing a closed queue."""


class EdgeQueueFull(EdgeQueueError):
    """Raised when a full queue uses the ``reject`` overflow policy."""


@dataclass(frozen=True, slots=True)
class QueueStats:
    """Point-in-time counters for one edge queue."""

    capacity: int
    current_size: int
    high_watermark: int
    enqueued: int
    dequeued: int
    dropped_oldest: int
    dropped_newest: int
    rejected: int
    blocked_puts: int
    discarded_on_close: int
    closed: bool


class EdgeQueue(Generic[ItemT]):
    """A small asyncio queue whose overload policy is never implicit.

    Unlike :class:`asyncio.Queue`, closing wakes both blocked producers and
    consumers.  A normal close is drainable: existing items are returned before
    :class:`EdgeQueueClosed` is raised.  ``close(discard=True)`` is available for
    forced shutdown.
    """

    def __init__(
        self,
        capacity: int,
        overflow: OverflowPolicy | str = OverflowPolicy.BLOCK,
        *,
        name: str = "",
    ) -> None:
        if capacity < 1:
            raise ValueError("edge queue capacity must be at least 1")
        try:
            parsed_overflow = OverflowPolicy(overflow)
        except ValueError as exc:
            choices = ", ".join(policy.value for policy in OverflowPolicy)
            raise ValueError(
                f"unknown overflow policy {overflow!r}; use {choices}"
            ) from exc

        self._capacity = capacity
        self._overflow = parsed_overflow
        self._name = name
        self._items: Deque[ItemT] = deque()
        self._condition = asyncio.Condition()
        self._closed = False

        self._high_watermark = 0
        self._enqueued = 0
        self._dequeued = 0
        self._dropped_oldest = 0
        self._dropped_newest = 0
        self._rejected = 0
        self._blocked_puts = 0
        self._discarded_on_close = 0

    @property
    def name(self) -> str:
        return self._name

    @property
    def capacity(self) -> int:
        return self._capacity

    @property
    def overflow(self) -> OverflowPolicy:
        return self._overflow

    @property
    def closed(self) -> bool:
        return self._closed

    def qsize(self) -> int:
        return len(self._items)

    def empty(self) -> bool:
        return not self._items

    def full(self) -> bool:
        return len(self._items) >= self._capacity

    def stats(self) -> QueueStats:
        """Return a consistent-enough event-loop snapshot of queue counters."""

        return QueueStats(
            capacity=self._capacity,
            current_size=len(self._items),
            high_watermark=self._high_watermark,
            enqueued=self._enqueued,
            dequeued=self._dequeued,
            dropped_oldest=self._dropped_oldest,
            dropped_newest=self._dropped_newest,
            rejected=self._rejected,
            blocked_puts=self._blocked_puts,
            discarded_on_close=self._discarded_on_close,
            closed=self._closed,
        )

    async def put(self, item: ItemT) -> PutResult:
        """Put ``item`` according to this queue's overflow policy."""

        async with self._condition:
            self._raise_if_closed()

            result = PutResult.ENQUEUED
            if self._overflow is OverflowPolicy.BLOCK:
                if self.full():
                    self._blocked_puts += 1
                while self.full() and not self._closed:
                    await self._condition.wait()
                self._raise_if_closed()
            elif self.full() and self._overflow is OverflowPolicy.DROP_OLDEST:
                self._items.popleft()
                self._dropped_oldest += 1
                result = PutResult.DROPPED_OLDEST
            elif self.full() and self._overflow is OverflowPolicy.DROP_NEWEST:
                self._dropped_newest += 1
                return PutResult.DROPPED_NEWEST
            elif self.full() and self._overflow is OverflowPolicy.REJECT:
                self._rejected += 1
                queue_name = f" {self._name!r}" if self._name else ""
                raise EdgeQueueFull(f"edge queue{queue_name} is full")

            self._items.append(item)
            self._enqueued += 1
            self._high_watermark = max(self._high_watermark, len(self._items))
            self._condition.notify_all()
            return result

    async def get(self) -> ItemT:
        """Get the next item, waiting until data arrives or the queue closes."""

        async with self._condition:
            while not self._items and not self._closed:
                await self._condition.wait()
            if not self._items:
                raise EdgeQueueClosed(self._closed_message())

            item = self._items.popleft()
            self._dequeued += 1
            self._condition.notify_all()
            return item

    async def close(self, *, discard: bool = False) -> None:
        """Close the queue and wake all waiters.

        The method is idempotent.  By default queued items remain available to
        consumers.  With ``discard=True`` they are removed immediately and
        accounted for in ``discarded_on_close``.
        """

        async with self._condition:
            if discard and self._items:
                self._discarded_on_close += len(self._items)
                self._items.clear()
            self._closed = True
            self._condition.notify_all()

    async def __aiter__(self) -> AsyncIterator[ItemT]:
        while True:
            try:
                yield await self.get()
            except EdgeQueueClosed:
                return

    def _raise_if_closed(self) -> None:
        if self._closed:
            raise EdgeQueueClosed(self._closed_message())

    def _closed_message(self) -> str:
        if self._name:
            return f"edge queue {self._name!r} is closed"
        return "edge queue is closed"


__all__ = [
    "EdgeQueue",
    "EdgeQueueClosed",
    "EdgeQueueError",
    "EdgeQueueFull",
    "OverflowPolicy",
    "PutResult",
    "QueueStats",
]
