"""HandoffManager - Agent 间任务转交"""

import asyncio
import json
import logging
from typing import Any

from agentkit.core.protocol import HandoffMessage

logger = logging.getLogger(__name__)


class HandoffManager:
    """Handoff 管理器

    通过 Redis Pub/Sub 管理 Agent 间的任务转交。
    """

    def __init__(self, redis: Any = None, dispatcher: Any = None):
        self._redis = redis
        self._dispatcher = dispatcher
        self._handlers: dict[str, list[Any]] = {}

    def register_handler(self, agent_name: str, handler: Any) -> None:
        """注册 Handoff 处理器"""
        if agent_name not in self._handlers:
            self._handlers[agent_name] = []
        self._handlers[agent_name].append(handler)

    async def send_handoff(self, handoff: HandoffMessage) -> None:
        """发送 Handoff 请求"""
        if self._redis is None:
            raise RuntimeError("Handoff requires Redis connection")

        channel = f"agent:{handoff.target_agent}:handoff"
        await self._redis.publish(channel, json.dumps(handoff.to_dict()))
        logger.info(
            f"Handoff sent: {handoff.source_agent} -> {handoff.target_agent} "
            f"(task={handoff.task_id}, reason={handoff.reason})"
        )

    async def listen_for_handoffs(self, agent_name: str) -> None:
        """监听发往指定 Agent 的 Handoff 请求"""
        if self._redis is None:
            return

        channel = f"agent:{agent_name}:handoff"
        pubsub = self._redis.pubsub()
        await pubsub.subscribe(channel)

        try:
            async for message in pubsub.listen():
                if message["type"] == "message":
                    data = json.loads(message["data"])
                    handoff = HandoffMessage.from_dict(data)
                    await self._handle_handoff(handoff)
        except asyncio.CancelledError:
            pass
        finally:
            await pubsub.unsubscribe(channel)

    async def _handle_handoff(self, handoff: HandoffMessage) -> None:
        """处理收到的 Handoff"""
        handlers = self._handlers.get(handoff.target_agent, [])
        for handler in handlers:
            try:
                await handler(handoff)
            except Exception as e:
                logger.error(f"Handoff handler error: {e}")
