"""HandoffTransport - Agent 间 Handoff 传输层抽象

提供统一的传输接口，支持：
- InProcessHandoffTransport: 进程内 asyncio.Queue 传输（Expert Team 模式）
- RedisHandoffTransport: Redis Pub/Sub 传输（分布式模式）
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
from typing import AsyncIterator, Awaitable, Callable

import redis.asyncio as aioredis

logger = logging.getLogger(__name__)


class HandoffTransport:
    """Handoff 传输层协议

    定义 Agent 间 Handoff 消息的传输接口，支持多种底层实现。
    """

    async def send(self, channel: str, message: dict) -> None:
        """发送消息到指定频道"""
        ...

    async def listen(self, channel: str) -> AsyncIterator[dict]:
        """监听指定频道的消息（异步生成器）"""
        ...
        yield

    def register_handler(self, channel: str, handler: Callable[[dict], Awaitable[None]]) -> None:
        """注册消息处理器"""
        ...


class InProcessHandoffTransport:
    """进程内 Handoff 传输

    基于 asyncio.Queue 实现的进程内消息传输，适用于 Expert Team 模式下
    同一进程内的 Agent 间 Handoff，无需 Redis 依赖。

    支持广播模式：同一频道可有多个消费者，每条消息会投递到所有消费者。
    """

    _CLOSED_SENTINEL: dict = {"_closed": True}
    _MAX_QUEUE_SIZE: int = 1024

    def __init__(self) -> None:
        self._channels: dict[str, list[asyncio.Queue[dict]]] = {}
        self._handlers: dict[str, list[Callable[[dict], Awaitable[None]]]] = {}
        self._closed = False

    async def send(self, channel: str, message: dict) -> None:
        """发送消息到指定频道

        将消息放入该频道所有消费者的队列中，并调用所有已注册的处理器。
        如果频道不存在（无消费者），消息将被丢弃。
        """
        queues = self._channels.get(channel, [])
        for queue in queues:
            await queue.put(message)

        handlers = self._handlers.get(channel, [])
        for handler in handlers:
            try:
                await handler(message)
            except Exception as e:
                logger.error(f"InProcessHandoffTransport handler error on channel '{channel}': {e}")

    async def listen(self, channel: str) -> AsyncIterator[dict]:  # type: ignore[override]
        """监听指定频道的消息

        为调用者创建一个独立的队列，通过异步生成器持续产出消息。
        每个调用者获得独立的队列，实现广播语义。
        """
        queue: asyncio.Queue[dict] = asyncio.Queue(maxsize=self._MAX_QUEUE_SIZE)
        if channel not in self._channels:
            self._channels[channel] = []
        self._channels[channel].append(queue)

        try:
            while True:
                message = await queue.get()
                if message is self._CLOSED_SENTINEL:
                    break
                yield message
        finally:
            # 清理：移除当前消费者的队列
            if channel in self._channels:
                try:
                    self._channels[channel].remove(queue)
                except ValueError:
                    pass
                if not self._channels[channel]:
                    del self._channels[channel]

    def register_handler(self, channel: str, handler: Callable[[dict], Awaitable[None]]) -> None:
        """注册消息处理器

        当频道收到消息时，处理器将被调用。
        """
        if channel not in self._handlers:
            self._handlers[channel] = []
        self._handlers[channel].append(handler)

    def close(self) -> None:
        """关闭传输，清理所有频道和处理器。

        向所有活跃的监听器队列放入哨兵值，使其能够优雅退出。
        """
        self._closed = True
        # Put sentinel into every queue so active listeners unblock and exit
        for channel_queues in self._channels.values():
            for queue in channel_queues:
                try:
                    queue.put_nowait(self._CLOSED_SENTINEL)
                except asyncio.QueueFull:
                    pass
        self._channels.clear()
        self._handlers.clear()


class RedisHandoffTransport:
    """Redis Pub/Sub Handoff 传输

    基于 Redis Pub/Sub 实现的分布式消息传输，适用于跨进程、跨节点的
    Agent 间 Handoff。

    使用懒连接：在首次使用时才建立 Redis 连接，而非初始化时。
    """

    def __init__(self, redis_url: str) -> None:
        self._redis_url = redis_url
        self._redis: aioredis.Redis | None = None
        self._active_pubsubs: list[aioredis.client.PubSub] = []
        self._handlers: dict[str, list[Callable[[dict], Awaitable[None]]]] = {}
        self._listen_tasks: dict[str, asyncio.Task] = {}

    async def _ensure_connection(self) -> aioredis.Redis:
        """确保 Redis 连接已建立（懒初始化）"""
        if self._redis is None:
            try:
                # DEPLOY-2: 显式传入 REDIS_PASSWORD env var（Helm 注入但 redis-py 不自动读取）。
                self._redis = aioredis.from_url(
                    self._redis_url,
                    decode_responses=True,
                    password=os.environ.get("REDIS_PASSWORD") or None,
                )
                await self._redis.ping()
                logger.info("RedisHandoffTransport connected to Redis")
            except Exception:
                # Reset on failure so future calls retry
                self._redis = None
                raise
        return self._redis

    async def send(self, channel: str, message: dict) -> None:
        """发送消息到指定频道

        通过 Redis PUBLISH 命令将消息发布到频道。
        """
        redis = await self._ensure_connection()
        await redis.publish(channel, json.dumps(message))
        logger.debug(f"RedisHandoffTransport sent message to channel '{channel}'")

    async def listen(self, channel: str) -> AsyncIterator[dict]:  # type: ignore[override]
        """监听指定频道的消息

        订阅 Redis 频道，通过异步生成器持续产出消息。
        """
        redis = await self._ensure_connection()
        pubsub = redis.pubsub()
        await pubsub.subscribe(channel)
        self._active_pubsubs.append(pubsub)

        try:
            async for message in pubsub.listen():
                if message["type"] == "message":
                    data = message["data"]
                    if isinstance(data, str):
                        yield json.loads(data)
                    else:
                        yield json.loads(data.decode())
        except asyncio.CancelledError:
            pass
        finally:
            await pubsub.unsubscribe(channel)
            await pubsub.aclose()
            if pubsub in self._active_pubsubs:
                self._active_pubsubs.remove(pubsub)

    def register_handler(self, channel: str, handler: Callable[[dict], Awaitable[None]]) -> None:
        """注册消息处理器并启动后台监听任务

        注册处理器后，自动启动一个后台 asyncio.Task 来监听该频道的消息。
        """
        if channel not in self._handlers:
            self._handlers[channel] = []
        self._handlers[channel].append(handler)

        # 启动后台监听任务
        if channel not in self._listen_tasks:
            self._listen_tasks[channel] = asyncio.create_task(self._background_listen(channel))

    async def _background_listen(self, channel: str) -> None:
        """后台监听频道消息并调用处理器"""
        try:
            async for message in self.listen(channel):
                handlers = self._handlers.get(channel, [])
                for handler in handlers:
                    try:
                        await handler(message)
                    except Exception as e:
                        logger.error(
                            f"RedisHandoffTransport handler error on channel '{channel}': {e}"
                        )
        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.error(f"RedisHandoffTransport listen error on channel '{channel}': {e}")

    async def close(self) -> None:
        """关闭传输，取消所有监听任务并关闭 Redis 连接"""
        # 取消所有后台监听任务
        for task in self._listen_tasks.values():
            if not task.done():
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
        self._listen_tasks.clear()
        self._handlers.clear()

        # 关闭所有活跃的 PubSub 连接
        for pubsub in self._active_pubsubs:
            try:
                await pubsub.unsubscribe()
                await pubsub.aclose()
            except Exception:
                pass
        self._active_pubsubs.clear()

        # 关闭 Redis 连接
        if self._redis is not None:
            await self._redis.close()
            self._redis = None