# 中文注释：
# 文件：echobot/agent/multi.py
# 说明：多 Agent 管理器，负责协调多个 Agent 实例。

"""Multi-Agent Manager."""

import asyncio
import uuid
from pathlib import Path
from typing import Any

from loguru import logger

from echobot.agent.instance import AgentInstance
from echobot.bus.events import InboundMessage, OutboundMessage
from echobot.bus.queue import MessageBus
from echobot.channels.telegram import TelegramChannel
from echobot.config.schema import (
    MultiAgentsConfig,
    AgentInstanceConfig,
    AgentBinding,
    CollaborationConfig,
    TelegramConfig,
)
from echobot.providers.litellm_provider import LiteLLMProvider


class MultiAgentManager:
    """
    多 Agent 管理器

    负责：
    - 加载和管理多个 Agent 实例
    - 消息路由（根据 bindings 规则）
    - Agent 间协作
    - 统一的 Telegram Channel（共享同一个 Bot Token）
    """

    def __init__(
        self,
        config: MultiAgentsConfig,
        provider: LiteLLMProvider,
        defaults: dict[str, Any],
        mcp_config: dict | None = None,
    ):
        """
        Initialize Multi-Agent Manager.

        Args:
            config: 多 Agent 配置
            provider: LLM Provider
            defaults: 全局默认配置
            mcp_config: MCP 配置
        """
        self.config = config
        self.provider = provider
        self.defaults = defaults
        self.mcp_config = mcp_config

        # Agent 实例字典
        self.agents: dict[str, AgentInstance] = {}

        # 协作配置
        self._collaboration_enabled = config.collaboration.enabled
        self._allow_calls = set(config.collaboration.allow_calls)

        # 待处理的 Agent 间调用结果
        self._pending_calls: dict[str, asyncio.Future] = {}

        # 共享的 Message Bus 和 Telegram Channel（所有 Agent 使用同一个 Bot Token）
        self._shared_bus: MessageBus | None = None
        self._shared_channel: TelegramChannel | None = None
        self._primary_token: str | None = None

    def load_agents(self) -> None:
        """加载所有 enabled 的 Agent 实例"""
        # 检查是否所有 Agent 使用同一个 Token
        tokens = set(a.token for a in self.config.agents if a.enabled and a.token)
        if len(tokens) == 1:
            # 所有 Agent 使用同一个 Token，使用共享的 Channel 和 Bus
            self._primary_token = list(tokens)[0]
            self._shared_bus = MessageBus()
            tg_config = TelegramConfig(enabled=True, token=self._primary_token)
            self._shared_channel = TelegramChannel(config=tg_config, bus=self._shared_bus)
            logger.info(f"All agents using same token, using shared Telegram channel")

        # 加载 Agent 实例
        for agent_config in self.config.agents:
            if not agent_config.enabled:
                logger.info(f"Agent {agent_config.id} is disabled, skipping")
                continue

            # 检查 token
            if not agent_config.token:
                logger.warning(f"Agent {agent_config.id} has no token, skipping")
                continue

            # 如果使用共享 channel，传递给它
            channel = self._shared_channel if self._shared_channel else None
            bus = self._shared_bus if self._shared_bus else None

            # 创建 Agent 实例
            agent = AgentInstance(
                config=agent_config,
                provider=self.provider,
                defaults=self.defaults,
                collaboration_config={
                    "enabled": self._collaboration_enabled,
                    "allow_calls": list(self._allow_calls),
                },
                shared_channel=channel,
                shared_bus=bus,
                mcp_config=self.mcp_config,
            )

            self.agents[agent_config.id] = agent
            logger.info(f"Loaded agent: {agent_config.id} ({agent_config.name})")

    async def start_all(self) -> None:
        """启动所有 Agent 实例"""
        logger.info(f"Starting {len(self.agents)} agents...")

        # 设置出站消息回调和协作回调
        for agent in self.agents.values():
            agent.set_outbound_callback(self._forward_outbound)
            agent.set_collaboration_callbacks(
                call_agent_callback=self.call_agent,
                get_available_agents_callback=self.get_available_agents,
            )

        # 如果使用共享 channel，需要设置消息路由器
        if self._shared_channel and self._shared_bus:
            # 启动共享的 Telegram Channel
            asyncio.create_task(self._shared_channel.start())
            # 启动消息路由任务
            asyncio.create_task(self._route_messages())
            # 启动 Agent 时不启动各自的 channel
            start_tasks = [agent.start(start_channel=False) for agent in self.agents.values()]
        else:
            # 各自启动自己的 channel
            start_tasks = [agent.start() for agent in self.agents.values()]

        await asyncio.gather(*start_tasks, return_exceptions=True)

        logger.info(f"All {len(self.agents)} agents started")

    async def _route_messages(self) -> None:
        """从共享 bus 中路由消息到对应的 Agent"""
        logger.info("Starting message router for shared channel")

        while True:
            try:
                msg = await asyncio.wait_for(self._shared_bus.consume_inbound(), timeout=1.0)
                if msg:
                    # 根据 mention 路由消息
                    agent = await self.route_message(msg)
                    if agent:
                        # 将消息发送到目标 Agent 的 bus，而不是共享 bus
                        await agent.bus.publish_inbound(msg)
                        logger.info(f"Routed message to agent: {agent.id}")
                    else:
                        logger.warning(f"Message could not be routed: {msg.content[:50]}...")
            except asyncio.TimeoutError:
                continue
            except Exception as e:
                logger.error(f"Error routing message: {e}")

    async def stop_all(self) -> None:
        """停止所有 Agent 实例"""
        logger.info(f"Stopping {len(self.agents)} agents...")

        # 停止所有 Agent
        stop_tasks = [agent.stop() for agent in self.agents.values()]
        await asyncio.gather(*stop_tasks, return_exceptions=True)

        # 停止共享的 Telegram Channel
        if self._shared_channel:
            await self._shared_channel.stop()

        logger.info("All agents stopped")

    async def _forward_outbound(self, msg: OutboundMessage) -> None:
        """
        转发出站消息

        这里可以添加全局的出站消息处理逻辑
        """
        # 可以在这里添加日志、统计等
        logger.debug(f"Outbound from {msg.channel}:{msg.chat_id}: {msg.content[:50]}...")

    def get_agent(self, agent_id: str) -> AgentInstance | None:
        """获取指定 Agent 实例"""
        return self.agents.get(agent_id)

    def list_agents(self) -> list[dict]:
        """列出所有 Agent 状态"""
        return [agent.get_status() for agent in self.agents.values()]

    async def route_message(self, msg: InboundMessage) -> AgentInstance | None:
        """
        路由消息到正确的 Agent

        路由优先级：
        1. @mention 调用
        2. bindings 配置规则
        3. 默认路由到第一个 Agent

        Args:
            msg: 入站消息

        Returns:
            目标 Agent 实例，如果没有匹配的则返回 None
        """
        # 1. 检查是否是 @mention 调用
        mention = msg.metadata.get("mention")
        if mention:
            # 尝试匹配 agent_id 或 name
            if mention in self.agents:
                logger.info(f"Routing message to agent (by mention id): {mention}")
                return self.agents[mention]

            # 尝试通过 name 匹配
            for agent in self.agents.values():
                if agent.name.lower() == mention.lower():
                    logger.info(f"Routing message to agent (by mention name): {agent.id}")
                    return agent

        # 2. 检查 bindings 配置
        # 首先检查是否有 broadcast 模式
        broadcast_agents = []
        for binding in self.config.bindings:
            if self._match_binding(msg, binding):
                if binding.broadcast:
                    agent = self.agents.get(binding.agent_id)
                    if agent:
                        broadcast_agents.append(agent)
                        logger.info(f"Broadcast to agent: {binding.agent_id}")

        if broadcast_agents:
            # 返回第一个 agent，但所有 agent 都会被通知
            # 这里暂时返回列表的第一个，后续需要修改调用方
            self._broadcast_to_agents(broadcast_agents, msg)
            return broadcast_agents[0] if broadcast_agents else None

        # 非广播模式：只路由到第一个匹配的 agent
        for binding in self.config.bindings:
            if self._match_binding(msg, binding):
                agent = self.agents.get(binding.agent_id)
                if agent:
                    logger.info(f"Routing message to agent (by binding): {binding.agent_id}")
                    return agent

        # 3. 默认路由到第一个可用的 Agent
        if self.agents:
            default_agent = next(iter(self.agents.values()))
            logger.info(f"Routing message to default agent: {default_agent.id}")
            return default_agent

        logger.warning("No agent available for routing")
        return None

    def _match_binding(self, msg: InboundMessage, binding: AgentBinding) -> bool:
        """检查消息是否匹配绑定规则"""
        match = binding.match

        # 检查 channel
        if msg.channel != match.channel:
            return False

        # 检查 chat_type
        msg_chat_type = msg.metadata.get("chat_type")
        if match.chat_type and msg_chat_type != match.chat_type:
            # 兼容 supergroup 和 group
            if match.chat_type == "group" and msg_chat_type not in ("group", "supergroup"):
                return False

        # 检查 mention
        if match.mention:
            msg_mention = msg.metadata.get("mention", "").lower()
            if msg_mention != match.mention.lower():
                return False

        return True

    async def _broadcast_to_agents(self, agents: list, msg: InboundMessage) -> None:
        """将消息广播给多个 Agent"""
        import asyncio

        # 为每个 agent 创建独立的消息副本
        for agent in agents:
            # 创建消息副本，避免并发问题
            msg_copy = InboundMessage(
                channel=msg.channel,
                sender_id=msg.sender_id,
                chat_id=msg.chat_id,
                content=msg.content,
                media=msg.media.copy() if msg.media else [],
                metadata=msg.metadata.copy() if msg.metadata else {},
            )
            # 发布到 agent 的消息总线
            asyncio.create_task(agent.bus.publish_inbound(msg_copy))
            logger.info(f"Broadcast message to agent: {agent.id}")

    # =========================================================================
    # Agent 间调用
    # =========================================================================

    async def call_agent(
        self,
        caller_id: str,
        target_agent_id: str,
        task: str,
        context: dict | None = None,
        timeout: float = 120.0,
    ) -> str:
        """
        调用另一个 Agent

        Args:
            caller_id: 调用者 Agent ID
            target_agent_id: 目标 Agent ID
            task: 任务描述
            context: 上下文信息
            timeout: 超时时间（秒）

        Returns:
            目标 Agent 的响应
        """
        # 1. 检查协作是否启用
        if not self._collaboration_enabled:
            return "Error: Agent collaboration is disabled"

        # 2. 检查权限
        if target_agent_id not in self._allow_calls:
            return f"Error: Agent '{target_agent_id}' is not allowed to be called"

        # 3. 获取目标 Agent
        target_agent = self.get_agent(target_agent_id)
        if not target_agent:
            return f"Error: Agent '{target_agent_id}' not found"

        # 4. 创建调用 ID
        call_id = str(uuid.uuid4())

        # 5. 创建结果 Future
        result_future: asyncio.Future[str] = asyncio.Future()
        self._pending_calls[call_id] = result_future

        try:
            # 6. 发送任务到目标 Agent
            msg = InboundMessage(
                channel="internal",
                sender_id=caller_id,
                chat_id=f"internal:{call_id}",
                content=task,
                metadata={
                    "call_id": call_id,
                    "caller_id": caller_id,
                    **(context or {}),
                },
            )

            await target_agent.bus.publish_inbound(msg)

            # 7. 等待结果
            result = await asyncio.wait_for(result_future, timeout=timeout)
            return result

        except asyncio.TimeoutError:
            return f"Error: Agent '{target_agent_id}' call timed out after {timeout}s"
        except Exception as e:
            return f"Error calling agent '{target_agent_id}': {str(e)}"
        finally:
            self._pending_calls.pop(call_id, None)

    async def _handle_agent_response(self, call_id: str, response: str) -> None:
        """处理 Agent 响应（内部方法）"""
        future = self._pending_calls.get(call_id)
        if future and not future.done():
            future.set_result(response)

    def get_available_agents(self) -> list[dict]:
        """获取可被调用的 Agent 列表"""
        return [
            {"id": agent.id, "name": agent.name, "role": agent.role}
            for agent in self.agents.values()
            if agent.id in self._allow_calls
        ]
