# SPDX-FileCopyrightText: © 2026 scy
#
# SPDX-License-Identifier: MIT

import asyncio
from asyncio import Queue
from collections import defaultdict
from collections.abc import AsyncGenerator, Generator, Iterable
from contextlib import contextmanager
from enum import Enum
from typing import overload


class IfFull(Enum):
    """Behavior to use when a subscriber's queue is full."""

    SKIP = 1
    WAIT = 2
    FAIL = 3


class AsyncQueueBus[T]:
    """Event bus based on Python asyncio queues."""

    def __init__(self, *, queue_size: int, if_full: IfFull = IfFull.WAIT):
        self._queue_size = queue_size
        self._if_full = if_full

        self._subscribers: defaultdict[str, list[Queue[T]]] = defaultdict(list)

    @staticmethod
    async def _get(topic: str, queue: Queue[T]) -> tuple[str, T]:
        return topic, (await queue.get())

    @overload
    @classmethod
    def _multi_consume(
        cls, queues: dict[str, Queue[T]], *, timeout: None = None
    ) -> AsyncGenerator[tuple[str, T]]: ...

    @overload
    @classmethod
    def _multi_consume(
        cls, queues: dict[str, Queue[T]], *, timeout: float | int = 0
    ) -> AsyncGenerator[tuple[str, T] | tuple[None, None]]: ...

    @classmethod
    async def _multi_consume(
        cls,
        queues: dict[str, Queue[T]],
        *,
        timeout: float | int | None = None,  # noqa: ASYNC109
    ) -> AsyncGenerator[tuple[str, T] | tuple[None, None]]:
        # For each queue, create a task that tries to retrieve an item.
        tasks = {
            asyncio.create_task(cls._get(topic, queue))
            for topic, queue in queues.items()
        }
        try:
            while True:
                # Wait for the first (or multiple) queues that have an item.
                done, _pending = await asyncio.wait(
                    tasks, timeout=timeout, return_when=asyncio.FIRST_COMPLETED
                )
                # If no task is done, we had a timeout. We signal that by
                # returning a `None` topic.
                if not done:
                    yield None, None
                    continue
                # For each of them, yield the topic and item, remove the
                # completed task, and schedule a new one on that queue.
                for task in done:
                    topic, value = await task
                    yield topic, value
                    tasks.remove(task)
                    tasks.add(
                        asyncio.create_task(cls._get(topic, queues[topic]))
                    )
        except asyncio.CancelledError:
            # Cancel all tasks and wait for them to be done.
            for task in tasks:
                task.cancel()
            await asyncio.wait(tasks)
            raise

    def _create_and_attach_queues(
        self, topics: Iterable[str], *, queue_size: int | None = None
    ) -> dict[str, Queue[T]]:
        # Fall back to instance's default queue size.
        queue_size = self._queue_size if queue_size is None else queue_size
        # Create queues for each of the topics.
        queues = {topic: Queue[T](queue_size) for topic in set(topics)}
        # Ensure all the topics are strings.
        for topic in queues:
            if not isinstance(topic, str):
                raise TypeError(
                    f"topics need to be strings, got {type(topic)}"
                )
        # Add each one to the subscriber lists.
        for topic, queue in queues.items():
            self._subscribers[topic].append(queue)
        return queues

    def _destroy_queues(self, queues: dict[str, Queue[T]]) -> None:
        for topic, queue in queues.items():
            self._subscribers[topic].remove(queue)
            # Clean up empty topics.
            if not self._subscribers[topic]:
                del self._subscribers[topic]

    async def publish(
        self, topic: str, data: T, if_full: IfFull | None = None
    ) -> None:
        """Publish an object to all subscribers to a topic."""
        if not isinstance(topic, str):
            raise TypeError(f"topic has to be a string, got {type(topic)}")
        if_full = if_full or self._if_full  # fall back to instance's default
        if topic not in self._subscribers:
            return  # nobody subscribes to that topic

        if if_full == IfFull.WAIT:
            # Create tasks for all subscribers, then wait for all of them.
            async with asyncio.TaskGroup() as tg:
                for queue in self._subscribers[topic]:
                    tg.create_task(queue.put(data))
            return

        self.sync_publish(topic, data, fail_if_full=if_full == IfFull.FAIL)

    @overload
    @contextmanager
    def subscribe(
        self, *topics: str, queue_size: int | None = None, timeout: None = None
    ) -> Generator[AsyncGenerator[tuple[str, T]]]: ...

    @overload
    @contextmanager
    def subscribe(
        self,
        *topics: str,
        queue_size: int | None = None,
        timeout: float | int = 0,
    ) -> Generator[AsyncGenerator[tuple[str, T] | tuple[None, None]]]: ...

    @contextmanager
    def subscribe(
        self,
        *topics: str,
        queue_size: int | None = None,
        timeout: float | int | None = None,
    ) -> Generator[AsyncGenerator[tuple[str, T] | tuple[None, None]]]:
        """Subscribe to one or more topics."""
        queues = self._create_and_attach_queues(topics, queue_size=queue_size)
        try:
            # The context manager's value is this function that returns tuples
            # of topics and values.
            yield self._multi_consume(queues, timeout=timeout)
        finally:
            # When the context manager exits, remove our queues.
            self._destroy_queues(queues)

    @contextmanager
    def subscribe_queue(
        self, topic: str, queue_size: int | None = None
    ) -> Generator[Queue[T]]:
        """Subscribe to exactly one topic and get its queue back."""
        queues = self._create_and_attach_queues({topic}, queue_size=queue_size)
        try:
            # The context manager's value is a dict with the queues.
            yield queues[topic]
        finally:
            # When the context manager exists, remove the queues again.
            self._destroy_queues(queues)

    @contextmanager
    def subscribe_queues(
        self, *topics: str, queue_size: int | None = None
    ) -> Generator[dict[str, Queue[T]]]:
        """Subscribe to one or more topics and get the raw queues back."""
        queues = self._create_and_attach_queues(topics, queue_size=queue_size)
        try:
            # The context manager's value is a dict with the queues.
            yield queues
        finally:
            # When the context manager exists, remove the queues again.
            self._destroy_queues(queues)

    def subscriber_count(self) -> dict[str, int]:
        """Return the number of subscribers for each known topic."""
        return {
            topic: len(queues) for topic, queues in self._subscribers.items()
        }

    def sync_publish(
        self, topic: str, data: T, *, fail_if_full: bool = True
    ) -> None:
        """Publish an object to all subscribers to a topic, without async.

        In contrast to publish(), you cannot wait for full queues, only skip or
        ignore them.
        """
        if not isinstance(topic, str):
            raise TypeError(f"topic has to be a string, got {type(topic)}")
        # Fill non-full queues first.
        full: Queue[T] | None = None
        for queue in self._subscribers[topic]:
            if queue.full():
                full = queue
            else:
                queue.put_nowait(data)

        if fail_if_full and full:
            # Put into the last full queue we saw, causing an exception.
            full.put_nowait(data)
        # Else we're simply skipping the full queues, nothing left to do.
