"""Async message bus for inter-agent communication."""

from __future__ import annotations

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

from vastian.bus.message import Message, MessageType

logger = logging.getLogger(__name__)

# Type alias for message handlers
MessageHandler = Callable[[Message], Awaitable[None]]


class MessageBus:
    """
    In-process async message bus.

    Supports:
    - Direct agent-to-agent messaging
    - Team broadcasts (recipients starting with "team:")
    - Global broadcasts (recipient = "broadcast")
    - Typed subscriptions (subscribe to specific MessageTypes)
    - Message history for audit trail
    """

    def __init__(self, max_history: int = 10_000) -> None:
        # Per-agent mailboxes
        self._mailboxes: dict[str, asyncio.Queue[Message]] = {}
        # Type-based subscribers: MessageType -> list of handler coroutines
        self._type_subscribers: dict[MessageType, list[MessageHandler]] = defaultdict(list)
        # Global subscribers (receive everything)
        self._global_subscribers: list[MessageHandler] = []
        # Team membership: team_id -> set of agent_ids
        self._teams: dict[str, set[str]] = defaultdict(set)
        # Audit trail
        self._history: list[Message] = []
        self._max_history = max_history
        self._lock = asyncio.Lock()

    # ── Registration ─────────────────────────────────────

    def register_agent(self, agent_id: str) -> asyncio.Queue[Message]:
        """Register an agent and return its mailbox queue."""
        if agent_id not in self._mailboxes:
            self._mailboxes[agent_id] = asyncio.Queue()
            logger.info("Registered agent on bus: %s", agent_id)
        return self._mailboxes[agent_id]

    def register_team(self, team_id: str, member_ids: list[str]) -> None:
        """Register a team with its member agent IDs."""
        self._teams[team_id] = set(member_ids)
        logger.info("Registered team %s with members: %s", team_id, member_ids)

    def subscribe(self, msg_type: MessageType, handler: MessageHandler) -> None:
        """Subscribe a handler to a specific message type."""
        self._type_subscribers[msg_type].append(handler)

    def subscribe_all(self, handler: MessageHandler) -> None:
        """Subscribe a handler to ALL messages (for logging, audit, etc.)."""
        self._global_subscribers.append(handler)

    # ── Sending ──────────────────────────────────────────

    async def send(self, message: Message) -> None:
        """Send a message to its recipient(s)."""
        async with self._lock:
            self._record(message)

        # Route based on recipient
        recipient = message.recipient

        if recipient == "broadcast":
            await self._broadcast(message)
        elif recipient.startswith("team:"):
            team_id = recipient.removeprefix("team:")
            await self._send_to_team(team_id, message)
        else:
            await self._send_to_agent(recipient, message)

        # Notify type-based subscribers
        for handler in self._type_subscribers.get(message.type, []):
            try:
                await handler(message)
            except Exception:
                logger.exception("Type subscriber error for %s", message.type)

        # Notify global subscribers
        for handler in self._global_subscribers:
            try:
                await handler(message)
            except Exception:
                logger.exception("Global subscriber error")

    async def _send_to_agent(self, agent_id: str, message: Message) -> None:
        """Deliver a message to a specific agent's mailbox."""
        mailbox = self._mailboxes.get(agent_id)
        if mailbox is None:
            logger.warning("No mailbox for agent %s — message dropped: %s", agent_id, message.id)
            return
        await mailbox.put(message)
        logger.debug("Delivered %s → %s [%s]", message.sender, agent_id, message.type.value)

    async def _send_to_team(self, team_id: str, message: Message) -> None:
        """Broadcast a message to all members of a team."""
        members = self._teams.get(team_id, set())
        if not members:
            logger.warning("Unknown team %s — message dropped: %s", team_id, message.id)
            return
        for member_id in members:
            if member_id != message.sender:  # Don't send to self
                await self._send_to_agent(member_id, message)

    async def _broadcast(self, message: Message) -> None:
        """Broadcast a message to all registered agents."""
        for agent_id in self._mailboxes:
            if agent_id != message.sender:
                await self._send_to_agent(agent_id, message)

    # ── History ──────────────────────────────────────────

    def _record(self, message: Message) -> None:
        """Append to audit trail, trimming if necessary."""
        self._history.append(message)
        if len(self._history) > self._max_history:
            self._history = self._history[-self._max_history :]

    def get_history(
        self,
        thread_id: str | None = None,
        sender: str | None = None,
        msg_type: MessageType | None = None,
        limit: int = 100,
    ) -> list[Message]:
        """Query message history with optional filters."""
        results = self._history
        if thread_id:
            results = [m for m in results if m.thread_id == thread_id]
        if sender:
            results = [m for m in results if m.sender == sender]
        if msg_type:
            results = [m for m in results if m.type == msg_type]
        return results[-limit:]

    @property
    def agent_ids(self) -> list[str]:
        return list(self._mailboxes.keys())

    @property
    def team_ids(self) -> list[str]:
        return list(self._teams.keys())
