"""RedisMessageBus — 基于 Redis Streams 的消息总线。

使用 XADD/XREADGROUP 实现可靠消息传递，支持消费者组、
消息确认和死信队列。
"""

from __future__ import annotations

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

from agentkit.bus.message import AgentMessage
from agentkit.bus.memory_bus import InMemoryMessageBus

logger = logging.getLogger(__name__)

_STREAM_PREFIX = "agentkit:bus:"
_DEAD_LETTER_SUFFIX = ":dead"


class RedisMessageBus:
    """基于 Redis Streams 的消息总线。"""

    def __init__(
        self,
        redis_url: str = "redis://localhost:6379/0",
        consumer_group: str = "agentkit_bus",
        max_retries: int = 3,
    ) -> None:
        self._redis_url = redis_url
        self._consumer_group = consumer_group
        self._max_retries = max_retries
        self._redis: Any = None
        self._subscribers: dict[str, list[Callable[[AgentMessage], Awaitable[None]]]] = {}
        self._pending_requests: dict[str, asyncio.Future[AgentMessage]] = {}
        self._consumer_tasks: dict[str, asyncio.Task] = {}

    async def _get_redis(self) -> Any:
        """获取 Redis 连接（懒初始化）。"""
        if self._redis is None:
            import redis.asyncio as aioredis
            self._redis = aioredis.from_url(self._redis_url, decode_responses=True)
        return self._redis

    def _stream_key(self, agent_name: str) -> str:
        return f"{_STREAM_PREFIX}{agent_name}"

    def _dead_letter_key(self, agent_name: str) -> str:
        return f"{_STREAM_PREFIX}{agent_name}{_DEAD_LETTER_SUFFIX}"

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

        redis = await self._get_redis()
        stream_key = self._stream_key(message.recipient)
        data = message.to_dict()

        try:
            await redis.xadd(stream_key, {"data": json.dumps(data)})
        except Exception as e:
            logger.error(f"Failed to publish message to {stream_key}: {e}")
            raise

        # Check pending requests (only for replies, not original request)
        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._subscribers[agent_name].append(handler)

        # Start consumer task
        if agent_name not in self._consumer_tasks:
            task = asyncio.create_task(
                self._consume_stream(agent_name),
            )
            self._consumer_tasks[agent_name] = task

    async def _consume_stream(self, agent_name: str) -> None:
        """消费 Redis Stream 中的消息。"""
        redis = await self._get_redis()
        stream_key = self._stream_key(agent_name)

        # Create consumer group if not exists
        try:
            await redis.xgroup_create(
                stream_key, self._consumer_group, id="0", mkstream=True,
            )
        except Exception:
            pass  # Group already exists

        while True:
            try:
                results = await redis.xreadgroup(
                    groupname=self._consumer_group,
                    consumername=agent_name,
                    streams={stream_key: ">"},
                    count=10,
                    block=1000,
                )

                if results:
                    for stream_name, messages in results:
                        for msg_id, fields in messages:
                            try:
                                data = json.loads(fields.get("data", "{}"))
                                message = AgentMessage.from_dict(data)

                                for handler in self._subscribers.get(agent_name, []):
                                    try:
                                        await handler(message)
                                    except Exception as e:
                                        logger.warning(f"Handler error for {agent_name}: {e}")

                                # Acknowledge message
                                await redis.xack(stream_key, self._consumer_group, msg_id)
                            except Exception as e:
                                logger.warning(f"Failed to process message {msg_id}: {e}")
                                # Move to dead letter after max retries
                                await self._handle_failed_message(
                                    redis, stream_key, msg_id, fields, agent_name,
                                )
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Consumer error for {agent_name}: {e}")
                await asyncio.sleep(1)

    async def _handle_failed_message(
        self,
        redis: Any,
        stream_key: str,
        msg_id: str,
        fields: dict,
        agent_name: str,
    ) -> None:
        """处理失败消息（移入死信队列）。"""
        dead_key = self._dead_letter_key(agent_name)
        try:
            await redis.xadd(dead_key, fields)
            await redis.xack(stream_key, self._consumer_group, msg_id)
            logger.warning(f"Message {msg_id} moved to dead letter queue")
        except Exception as e:
            logger.error(f"Failed to move message to dead letter: {e}")

    async def unsubscribe(self, agent_name: str) -> None:
        """取消订阅。"""
        self._subscribers.pop(agent_name, None)
        task = self._consumer_tasks.pop(agent_name, None)
        if task:
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass

    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:
        """广播消息。"""
        message.recipient = None

        redis = await self._get_redis()
        data = message.to_dict()

        for agent_name in self._subscribers:
            stream_key = self._stream_key(agent_name)
            try:
                await redis.xadd(stream_key, {"data": json.dumps(data)})
            except Exception as e:
                logger.error(f"Failed to broadcast to {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:
        try:
            redis = await self._get_redis()
            return await redis.ping()
        except Exception:
            return False

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


def create_message_bus(
    backend: str = "memory",
    redis_url: str = "redis://localhost:6379/0",
    consumer_group: str = "agentkit_bus",
    max_retries: int = 3,
) -> InMemoryMessageBus | RedisMessageBus:
    """创建消息总线实例。

    Args:
        backend: "memory" 或 "redis"
        redis_url: Redis 连接 URL
        consumer_group: Redis 消费者组名称
        max_retries: 消息最大重试次数

    Returns:
        MessageBus 实例
    """
    if backend == "redis":
        try:
            import redis.asyncio as aioredis  # noqa: F401
            bus = RedisMessageBus(
                redis_url=redis_url,
                consumer_group=consumer_group,
                max_retries=max_retries,
            )
            logger.info(f"MessageBus backend: redis_streams ({redis_url})")
            return bus
        except Exception as exc:
            logger.warning(
                f"Failed to initialise RedisMessageBus ({exc}), "
                f"falling back to InMemoryMessageBus"
            )

    bus = InMemoryMessageBus()
    logger.info("MessageBus backend: memory")
    return bus
