from __future__ import annotations

import asyncio
import inspect
import logging
from collections import defaultdict
from collections.abc import Awaitable, Callable
from typing import Any, TypeVar
from uuid import uuid4

from nexus.core.events.types import BaseEvent

logger = logging.getLogger(__name__)
T = TypeVar("T", bound=BaseEvent)


class EventBus:
    """Async pub/sub event bus using queue-backed subscriptions."""

    def __init__(self) -> None:
        self._subs: defaultdict[type[BaseEvent], list[tuple[str, asyncio.Queue[BaseEvent], Callable[[Any], Awaitable[None]]]]] = (
            defaultdict(list)
        )
        self._workers: dict[str, asyncio.Task[None]] = {}
        self._running = False

    def subscribe(self, event_type: type[T], handler: Callable[[T], Awaitable[None]]) -> str:
        """Register an async handler for an event type and return subscription id."""
        sub_id = str(uuid4())
        queue: asyncio.Queue[BaseEvent] = asyncio.Queue()
        self._subs[event_type].append((sub_id, queue, handler))
        if self._running:
            self._workers[sub_id] = asyncio.create_task(self._dispatch_worker(sub_id, queue, handler))
        return sub_id

    def unsubscribe(self, subscription_id: str) -> None:
        """Remove a subscription by id and stop its worker if running."""
        for event_type, entries in list(self._subs.items()):
            remaining: list[tuple[str, asyncio.Queue[BaseEvent], Callable[[Any], Awaitable[None]]]] = []
            for sub_id, queue, handler in entries:
                if sub_id == subscription_id:
                    task = self._workers.pop(sub_id, None)
                    if task and not task.done():
                        task.cancel()
                else:
                    remaining.append((sub_id, queue, handler))
            if remaining:
                self._subs[event_type] = remaining
            else:
                del self._subs[event_type]

    def publish(self, event: BaseEvent) -> None:
        """Publish event to matching subscribers without awaiting handlers."""
        for etype, subs in self._subs.items():
            if isinstance(event, etype):
                for _, queue, _ in subs:
                    try:
                        queue.put_nowait(event)
                    except asyncio.QueueFull:
                        logger.warning("Event dropped for %s due to full queue", etype.__name__)

    async def publish_and_wait(self, event: BaseEvent) -> None:
        """Publish and await all matching handlers directly."""
        tasks: list[asyncio.Task[None]] = []
        for etype, subs in self._subs.items():
            if isinstance(event, etype):
                for _, _, handler in subs:
                    tasks.append(asyncio.create_task(self._safe_handler_call(handler, event)))
        if tasks:
            await asyncio.gather(*tasks)

    async def start(self) -> None:
        """Start background queue workers for all current subscriptions."""
        if self._running:
            return
        self._running = True
        for _, entries in self._subs.items():
            for sub_id, queue, handler in entries:
                if sub_id not in self._workers or self._workers[sub_id].done():
                    self._workers[sub_id] = asyncio.create_task(self._dispatch_worker(sub_id, queue, handler))

    async def stop(self) -> None:
        """Drain pending events and stop workers."""
        self._running = False
        for _, entries in self._subs.items():
            for _, queue, _ in entries:
                while not queue.empty():
                    await asyncio.sleep(0)

        for task in self._workers.values():
            if not task.done():
                task.cancel()
        for task in list(self._workers.values()):
            try:
                await task
            except asyncio.CancelledError:
                continue
        self._workers.clear()

    async def _dispatch_worker(
        self,
        subscription_id: str,
        queue: asyncio.Queue[BaseEvent],
        handler: Callable[[Any], Awaitable[None]],
    ) -> None:
        """Drain subscriber queue and invoke its handler."""
        try:
            while True:
                event = await queue.get()
                await self._safe_handler_call(handler, event)
                queue.task_done()
        except asyncio.CancelledError:
            logger.debug("Subscription worker %s cancelled", subscription_id)
            raise

    async def _safe_handler_call(self, handler: Callable[[Any], Awaitable[None]], event: BaseEvent) -> None:
        try:
            maybe = handler(event)
            if inspect.isawaitable(maybe):
                await maybe
        except Exception as exc:
            logger.exception("Event handler failed for %s: %s", type(event).__name__, exc)
