"""
MessageBus — internal pub/sub communication backbone for Nexus agents.

Design notes:
  - In-process: agents in the same process communicate via async queues (zero-copy)
  - Supports topic-based pub/sub and point-to-point messaging
  - Transparent upgrade: same API whether agents are in-process or remote
  - Messages are typed (Message objects) for traceability
"""

from __future__ import annotations

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

from nexus.core.message import Message

logger = logging.getLogger("nexus.protocol.bus")

# Subscriber callback type
Subscriber = Callable[[Message], Awaitable[None]]


class MessageBus:
    """
    Async pub/sub message bus for intra-framework communication.

    Usage:
        bus = MessageBus()

        async def handler(msg: Message):
            print(f"Received: {msg.content}")

        bus.subscribe("research-results", handler)
        await bus.publish("research-results", Message.agent("Found 10 papers"))

        # Point-to-point
        await bus.send("agent-123", Message.user("Please analyze this"))
        msg = await bus.receive("agent-123")
    """

    def __init__(self) -> None:
        # Topic → list of subscribers
        self._subscribers: dict[str, list[Subscriber]] = defaultdict(list)
        # Agent-id → incoming message queue (point-to-point)
        self._queues: dict[str, asyncio.Queue[Message]] = {}
        # Message history (bounded, for debugging)
        self._history: list[Message] = []
        self._max_history = 1000
        self._lock = asyncio.Lock()

    # ---- Pub/Sub ----

    def subscribe(self, topic: str, callback: Subscriber) -> None:
        """Subscribe to messages on a topic."""
        self._subscribers[topic].append(callback)
        logger.debug(f"Subscriber added to topic '{topic}'")

    def unsubscribe(self, topic: str, callback: Subscriber) -> None:
        """Unsubscribe from a topic."""
        subs = self._subscribers.get(topic, [])
        if callback in subs:
            subs.remove(callback)

    async def publish(self, topic: str, message: Message) -> int:
        """
        Publish a message to all subscribers of a topic.

        Returns the number of subscribers notified.
        """
        self._record(message)
        subs = self._subscribers.get(topic, [])
        if not subs:
            logger.debug(f"No subscribers for topic '{topic}'")
            return 0

        # Fan out to all subscribers concurrently
        tasks = [sub(message) for sub in subs]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # Log any subscriber errors
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                logger.error(f"Subscriber {i} on topic '{topic}' failed: {result}")

        return len(subs)

    # ---- Point-to-Point ----

    async def send(self, agent_id: str, message: Message) -> None:
        """Send a message directly to an agent's queue."""
        self._record(message)
        message.receiver_id = agent_id
        queue = self._get_queue(agent_id)
        await queue.put(message)
        logger.debug(f"Message sent to agent '{agent_id}'")

    async def receive(
        self,
        agent_id: str,
        timeout: float | None = None,
    ) -> Message | None:
        """
        Receive the next message from an agent's queue.

        Returns None if timeout expires without a message.
        """
        queue = self._get_queue(agent_id)
        try:
            if timeout:
                return await asyncio.wait_for(queue.get(), timeout=timeout)
            return await queue.get()
        except TimeoutError:
            return None

    def has_pending(self, agent_id: str) -> bool:
        """Check if there are pending messages for an agent."""
        queue = self._queues.get(agent_id)
        return queue is not None and not queue.empty()

    # ---- Request/Response pattern ----

    async def request(
        self,
        target_agent_id: str,
        message: Message,
        timeout: float = 30.0,
    ) -> Message | None:
        """
        Send a message and wait for a response (request/response pattern).

        The response is expected to arrive on the sender's queue with a
        matching context_id.
        """
        sender_id = message.sender_id or "anonymous"
        context_id = message.context_id or message.id
        message.context_id = context_id

        await self.send(target_agent_id, message)

        # Wait for a response with matching context_id
        deadline = asyncio.get_event_loop().time() + timeout
        while asyncio.get_event_loop().time() < deadline:
            remaining = deadline - asyncio.get_event_loop().time()
            response = await self.receive(sender_id, timeout=min(remaining, 1.0))
            if response and response.context_id == context_id:
                return response

        logger.warning(f"Request to '{target_agent_id}' timed out after {timeout}s")
        return None

    # ---- Internals ----

    def _get_queue(self, agent_id: str) -> asyncio.Queue[Message]:
        if agent_id not in self._queues:
            self._queues[agent_id] = asyncio.Queue()
        return self._queues[agent_id]

    def _record(self, message: Message) -> None:
        self._history.append(message)
        if len(self._history) > self._max_history:
            self._history = self._history[-self._max_history:]

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

    @property
    def history(self) -> list[Message]:
        return list(self._history)

    def clear(self) -> None:
        """Reset the bus (useful for testing)."""
        self._subscribers.clear()
        self._queues.clear()
        self._history.clear()
