"""
Priority mapping library
"""

from __future__ import annotations

import anyio
from dataclasses import dataclass
from time import monotonic as time

from collections.abc import Hashable, MutableMapping
from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload

if TYPE_CHECKING:  # pragma: no cover
    from typing import Protocol

    class Comparable(Protocol):
        """Protocol for annotating comparable priorities."""

        def __lt__(self, other: Any, /) -> bool: ...

else:
    Comparable = Any

Priority = TypeVar("Priority", bound=Comparable)
KeyT = TypeVar("KeyT", bound=Hashable)

if TYPE_CHECKING:  # pragma: no cover
    from types import EllipsisType

    from collections.abc import (
        ItemsView,
        Iterable,
        Iterator,
        KeysView,
        ValuesView,
    )

    RT = TypeVar("RT")

    InitialData = dict[KeyT, Priority] | None
    InitialPrio = dict[KeyT, float] | None


@dataclass(slots=True)
class _KeyPrio(Generic[KeyT, Priority]):
    """Storage unit for heap entries."""

    key: KeyT
    priority: Priority


__all__ = ["PrioMap", "TimerMap"]


class PrioMap(MutableMapping, Generic[KeyT, Priority]):
    """
    A heap that behaves like a dict but maintains heap ordering.

    Supports dictionary-like access, key updates,
    removals, bulk initialization, and safe iteration (detects concurrent modifications).
    """

    def __init__(self, initial: InitialData = None) -> None:
        """
        Initialize the HeapDict.

        Args:
            initial: Optional mapping of keys to initial priorities.
        Raises:
            TypeError: If any priority in ``initial`` is not an int or float.
        """
        self.heap: list[_KeyPrio[KeyT, Priority]] = []
        self.position: dict[KeyT, int] = {}
        self.evt: anyio.Event = anyio.Event()

        # Bulk initialize if provided
        if initial:
            self.bulk(initial.items())

    def bulk(self, initial: Iterable[tuple[KeyT, Priority]]) -> None:
        """
        Bulk insert.
        """
        for key, priority in initial:
            self.heap.append(_KeyPrio(key, priority))
        # Record positions and heapify
        for idx, item in enumerate(self.heap):
            self.position[item.key] = idx
        for i in reversed(range(len(self.heap) // 2)):
            self._sift_down(i)

    def items(self) -> ItemsView[KeyT, Priority]:
        """
        Yield (key, priority) pairs.

        Items are heap sorted, i.e. the first result is guaranteed to have
        the lowest priority, but after that it's anybody's guess.
        """
        return self._create_iterator(None)

    def keys(self) -> KeysView[KeyT]:
        """
        Yield keys only.

        Items are heap sorted, i.e. the first key is guaranteed to have
        the lowest priority, but after that it's anybody's guess.
        """
        return self._create_iterator(True)

    def values(self) -> ValuesView[Priority]:
        """
        Yield priorities only.

        Items are heap sorted, i.e. the first priority is guaranteed to be
        lowest, but after that it's anybody's guess.
        """
        return self._create_iterator(False)

    @overload
    def pop(self) -> tuple[KeyT, Priority]: ...

    @overload
    def pop(self, key: KeyT) -> Priority: ...

    @overload
    def pop(self, key: KeyT, default: RT) -> Priority | RT: ...

    def pop(self, *a) -> tuple[KeyT, Priority] | Priority | RT:
        """
        Remove and return an item.

        Args are passed to dict.pop.

        :return: (key, priority)
        :raises IndexError: If empty.
        """
        if a:
            if len(a) > 2:
                raise TypeError("PrioMap.pop([key[,default]])")
            try:
                pos = self.position[a[0]]
            except KeyError:
                if len(a) > 1:
                    return a[1]
                raise
        else:
            pos = 0

        item = self.heap[pos]
        key = item.key
        prio = item.priority

        last = self.heap.pop()
        if pos < len(self.heap):
            self.heap[pos] = last
            self.position[last.key] = pos
            self._sift_down(pos)
        del self.position[key]
        if a:
            return prio
        return key, prio

    def peek(self) -> tuple[KeyT, Priority]:
        """
        Return the root item without removing it.

        :raises IndexError: If empty.
        """
        try:
            item = self.heap[0]
            return item.key, item.priority
        except IndexError:
            raise IndexError("Queue is empty") from None

    def set_priority(self, key: KeyT, new_priority: Priority) -> None:
        """
        Update priority for an existing key, then reheapify.

        :param key: Key to update.
        :param new_priority: New priority value.
        :raises KeyError: If key not found.
        :raises TypeError: If new_priority invalid.
        """
        if key not in self.position:
            raise KeyError(f"Key {key} not found in heap.")
        idx = self.position[key]
        old = self.heap[idx].priority
        self.heap[idx].priority = new_priority
        if new_priority < old:
            self._sift_up(idx)
        else:
            self._sift_down(idx)

        if idx == 0 or self.heap[0].key == key:
            self.evt.set()
            self.evt = anyio.Event()

    def clear(self) -> None:
        """
        Remove all items from the heap.
        """
        self.heap.clear()
        self.position.clear()
        self.evt.set()
        self.evt = anyio.Event()

    def is_empty(self) -> bool:
        """
        Check whether heap is empty.

        :return: True if no items.
        """
        return not self.heap

    def __bool__(self) -> bool:
        """
        Check whether heap is not empty.

        :return: False if no items.
        """
        return bool(self.heap)

    def _swap(self, i: int, j: int) -> None:
        """
        Swap elements at indices `i` and `j` and update their positions.
        """
        self.heap[i], self.heap[j] = self.heap[j], self.heap[i]
        self.position[self.heap[i].key] = i
        self.position[self.heap[j].key] = j

    def _sift_up(self, idx: int) -> None:
        """
        Sift element at index `idx` up until heap property is restored.
        """
        while idx > 0:
            parent = (idx - 1) // 2
            if self.heap[idx].priority < self.heap[parent].priority:
                self._swap(idx, parent)
                idx = parent
            else:
                break

    def _sift_down(self, idx: int) -> None:
        """
        Sift element at index `idx` down until heap property is restored.
        """
        n = len(self.heap)
        while True:
            left = 2 * idx + 1
            right = 2 * idx + 2
            best = idx

            if left < n and self.heap[left].priority < self.heap[best].priority:
                best = left
            if right < n and self.heap[right].priority < self.heap[best].priority:
                best = right

            if best != idx:
                self._swap(idx, best)
                idx = best
            else:
                break

    def __getitem__(self, key: KeyT) -> Priority:
        """
        Get the priority for `key`.

        :param key: Key to look up.
        :return: Associated priority.
        :raises KeyError: If `key` is not present.
        """
        if key in self.position:
            return self.heap[self.position[key]].priority
        raise KeyError(f"Key {key} not found in heap.")

    def __setitem__(self, key: KeyT, priority: Priority) -> None:
        """
        Insert or update `key` with `priority`.

        :param key: Key to insert/update.
        :param priority: Priority value (int or float).
        :raises TypeError: If `priority` is not int or float.
        """
        if key in self.position:
            self.set_priority(key, priority)
        else:
            idx = len(self.heap)
            self.heap.append(_KeyPrio(key, priority))
            self.position[key] = idx
            self._sift_up(idx)
            if self.heap[0].key == key:
                self.evt.set()
                self.evt = anyio.Event()

    def __delitem__(self, key: KeyT) -> None:
        """
        Remove `key` from the heap.

        :param key: Key to remove.
        :raises KeyError: If `key` not present.
        """
        if key not in self.position:
            raise KeyError(f"Key {key} not found in heap.")
        idx = self.position.pop(key)
        last = self.heap.pop()
        if idx < len(self.heap):
            self.heap[idx] = last
            self.position[last.key] = idx
            self._sift_down(idx)
            self._sift_up(idx)
            if idx == 0:
                self.evt.set()
                self.evt = anyio.Event()

    def __contains__(self, key: object) -> bool:
        """
        Check if `key` exists in the heap.
        """
        return key in self.position

    def __len__(self) -> int:
        """
        Return number of items.
        """
        return len(self.heap)

    def __str__(self) -> str:
        """
        String representation: list of {key: priority}.
        """
        return "[" + ", ".join(f"{{{item.key}: {item.priority}}}" for item in self.heap) + "]"

    def _create_iterator(self, keys: bool | None = None) -> Any:
        """
        Internal: return iterator over keys, values, or items, detecting concurrent mods.

        :param keys: Yield keys if True, values if False, both if None
        """
        self._iterator_state = {
            "index": 0,
            "len": len(self.heap),
            "pos": self.position.copy(),
        }

        class SafeIterator:
            def __init__(self, heap_dict, keys) -> None:
                self.heap_dict = heap_dict
                self.state = heap_dict._iterator_state  # noqa: SLF001
                self.keys = keys

            def __iter__(self) -> SafeIterator:
                return self

            def __next__(self) -> KeyT | Priority | tuple[KeyT, Priority]:
                s = self.state
                if s["index"] < s["len"]:
                    item = self.heap_dict.heap[s["index"]]
                    s["index"] += 1
                    if s["pos"] != self.heap_dict.position:
                        raise RuntimeError("Modification detected during iteration.")
                    if self.keys:
                        return item.key
                    if self.keys is False:
                        return item.priority
                    return item.key, item.priority
                raise StopIteration

        return SafeIterator(self, keys)

    def __iter__(self) -> Iterator[tuple[KeyT, Priority]]:
        """
        Iterate over (key, priority) pairs.

        The contents are *not* consumed.
        """
        return self._create_iterator(None)

    def __aiter__(self) -> PrioMap[KeyT, Priority]:
        """
        Iterate asynchronously over (key, priority) pairs.

        Contents are consumed.
        """
        return self

    async def __anext__(self) -> tuple[KeyT, Priority]:
        """
        Return the lowest-priority item.

        Waits for the next item if the heap is empty.
        """
        while not self.heap:
            await self.evt.wait()
        return self.pop()

    async def apeek(self) -> tuple[KeyT, Priority]:
        """
        Return the root item without removing it.

        Waits for an item to arrive if the heap is empty.
        """
        while not self.heap:
            await self.evt.wait()
        item = self.heap[0]
        return item.key, item.priority


class TimerMap(Generic[KeyT]):
    """
    A map that stores timeout values.

    The value auto-decrements; async iteration returns a key
    when its timer expires.

    Timeouts can be retrieved and updated iff they have not been processed.
    Negative delays are not an error.
    """

    def __init__(self, initial: dict[KeyT, float] | None = None) -> None:
        self._pm: PrioMap[KeyT, float] = PrioMap()
        if initial:
            self._pm.bulk((k, self.T_ADD(v)) for k, v in initial.items())

    @staticmethod
    def T_ADD(p) -> float:
        "Add the current time."
        return p + time()

    @staticmethod
    def T_SUB(p) -> float:
        "Subtract the current time."
        return p - time()

    def __setitem__(self, key: KeyT, delay: float) -> None:
        self._pm[key] = self.T_ADD(delay)

    def __getitem__(self, key: KeyT) -> float:
        return self.T_SUB(self._pm[key])

    def __delitem__(self, key: KeyT) -> None:
        del self._pm[key]

    async def apeek(self) -> tuple[KeyT, float]:
        """
        Return the first item (without removing it).

        Waits while the heap is empty.
        """
        k, p = await self._pm.apeek()
        return k, self.T_SUB(p)

    def __len__(self) -> int:
        return len(self._pm)

    def __aiter__(self) -> TimerMap[KeyT]:
        return self

    async def __anext__(self) -> KeyT:
        "iterate keys as they time out."
        while True:
            k, p = await self.apeek()
            if p > 0:
                with anyio.move_on_after(p):
                    await self._pm.evt.wait()
            else:
                k, p = self.pop()
                if p > 0:
                    raise RuntimeError("Heap got confused? {k !r}:{p}")
                return k

    def update(self, key: KeyT, new_delay: float) -> None:
        """
        Update priority for an existing key, then reheapify.
        """
        self._pm.set_priority(key, self.T_ADD(new_delay))

    @overload
    def pop(self) -> tuple[KeyT, float]: ...

    @overload
    def pop(self, key: KeyT) -> float: ...

    def pop(self, a: KeyT | EllipsisType = Ellipsis) -> tuple[KeyT, float] | float:
        """
        Remove and return an item.

        Args are passed to dict.pop.

        :return: (key, priority)
        :raises IndexError: If empty.
        """
        if a is Ellipsis:
            key, prio = self._pm.pop()
            return key, self.T_SUB(prio)
        else:
            return self.T_SUB(self._pm.pop(a))
