"""InMemoryMessageBus — 基于 asyncio.Queue 的内存消息总线。

用于开发和测试，行为与 Redis 实现一致。
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any, Callable, Awaitable

from agentkit.bus.message import AgentMessage

logger = logging.getLogger(__name__)


class InMemoryMessageBus:
    """基于 asyncio.Queue 的内存消息总线。"""

    def __init__(self) -> None:
        self._subscribers: dict[str, list[Callable[[AgentMessage], Awaitable[None]]]] = {}
        self._pending_requests: dict[str, asyncio.Future[AgentMessage]] = {}
        self._queues: dict[str, asyncio.Queue[AgentMessage]] = {}

    async def publish(self, message: AgentMessage) -> None:
        """发布消息。"""
        if message.is_broadcast:
            await self.broadcast(message)
            return

        # Point-to-point: deliver to recipient's queue
        recipient = message.recipient
        if recipient and recipient in self._queues:
            await self._queues[recipient].put(message)
        elif recipient and recipient in self._subscribers:
            # No queue, call handlers directly
            for handler in self._subscribers[recipient]:
                try:
                    await handler(message)
                except Exception as e:
                    logger.warning(f"Handler error for {recipient}: {e}")

        # Check if this is a response to a pending request
        # Only resolve if this is a reply (message_id != correlation_id),
        # not the original request itself
        if (
            message.correlation_id
            and message.correlation_id in self._pending_requests
            and message.message_id != message.correlation_id
        ):
            future = self._pending_requests[message.correlation_id]
            if not future.done():
                future.set_result(message)

    async def subscribe(
        self,
        agent_name: str,
        handler: Callable[[AgentMessage], Awaitable[None]],
    ) -> None:
        """订阅消息。"""
        if agent_name not in self._subscribers:
            self._subscribers[agent_name] = []
            self._queues[agent_name] = asyncio.Queue()
        self._subscribers[agent_name].append(handler)

        # Start consumer task
        asyncio.create_task(self._consume_queue(agent_name, handler))

    async def _consume_queue(
        self,
        agent_name: str,
        handler: Callable[[AgentMessage], Awaitable[None]],
    ) -> None:
        """消费队列中的消息。"""
        queue = self._queues.get(agent_name)
        if queue is None:
            return
        while True:
            try:
                message = await queue.get()
                try:
                    await handler(message)
                except Exception as e:
                    logger.warning(f"Handler error for {agent_name}: {e}")
            except asyncio.CancelledError:
                break

    async def unsubscribe(self, agent_name: str) -> None:
        """取消订阅。"""
        self._subscribers.pop(agent_name, None)
        self._queues.pop(agent_name, None)

    async def request(
        self,
        message: AgentMessage,
        timeout: float = 30.0,
    ) -> AgentMessage:
        """请求-响应模式。"""
        if not message.correlation_id:
            message.correlation_id = message.message_id

        loop = asyncio.get_event_loop()
        future: asyncio.Future[AgentMessage] = loop.create_future()
        self._pending_requests[message.correlation_id] = future

        try:
            await self.publish(message)
            return await asyncio.wait_for(future, timeout=timeout)
        except asyncio.TimeoutError:
            raise TimeoutError(
                f"Request {message.correlation_id} timed out after {timeout}s"
            )
        finally:
            self._pending_requests.pop(message.correlation_id, None)

    async def broadcast(self, message: AgentMessage) -> None:
        """广播消息。"""
        # Ensure recipient is None for broadcast
        message.recipient = None

        for agent_name, handlers in self._subscribers.items():
            for handler in handlers:
                try:
                    await handler(message)
                except Exception as e:
                    logger.warning(f"Broadcast handler error for {agent_name}: {e}")

        # Check pending requests (only for replies)
        if (
            message.correlation_id
            and message.correlation_id in self._pending_requests
            and message.message_id != message.correlation_id
        ):
            future = self._pending_requests[message.correlation_id]
            if not future.done():
                future.set_result(message)

    async def health_check(self) -> bool:
        return True

    @property
    def backend_type(self) -> str:
        return "memory"