"""AgentTool - 将 Agent 包装为 Tool"""

from typing import Any

from agentkit.tools.base import Tool


class AgentTool(Tool):
    """将另一个 Agent 包装为 Tool

    通过 Dispatcher 分发任务到目标 Agent，等待结果返回。
    """

    def __init__(
        self,
        name: str,
        description: str,
        agent_name: str,
        task_type: str,
        input_mapping: dict[str, str] | None = None,
        output_mapping: dict[str, str] | None = None,
        timeout_seconds: int = 300,
        version: str = "1.0.0",
        tags: list[str] | None = None,
    ):
        super().__init__(
            name=name,
            description=description,
            version=version,
            tags=tags or ["agent"],
        )
        self.agent_name = agent_name
        self.task_type = task_type
        self.input_mapping = input_mapping or {}
        self.output_mapping = output_mapping or {}
        self.timeout_seconds = timeout_seconds
        self._dispatcher = None

    def set_dispatcher(self, dispatcher: Any) -> "AgentTool":
        """注入 Dispatcher"""
        self._dispatcher = dispatcher
        return self

    async def execute(self, **kwargs) -> dict:
        if self._dispatcher is None:
            raise RuntimeError(f"AgentTool '{self.name}' has no dispatcher configured")

        from agentkit.core.protocol import TaskMessage
        from datetime import datetime, timezone
        import uuid

        # 映射输入
        mapped_input = {}
        for target_key, source_key in self.input_mapping.items():
            if source_key in kwargs:
                mapped_input[target_key] = kwargs[source_key]
        if not mapped_input:
            mapped_input = kwargs

        task = TaskMessage(
            task_id=str(uuid.uuid4()),
            agent_name=self.agent_name,
            task_type=self.task_type,
            priority=0,
            input_data=mapped_input,
            callback_url=None,
            created_at=datetime.now(timezone.utc),
            timeout_seconds=self.timeout_seconds,
        )

        await self._dispatcher.dispatch(task)

        # 等待结果
        import asyncio
        for _ in range(self.timeout_seconds):
            status = await self._dispatcher.get_task_status(task.task_id)
            if status["status"] in ("completed", "failed", "cancelled"):
                if status["status"] == "completed" and status.get("output_data"):
                    output = status["output_data"]
                    # 映射输出
                    if self.output_mapping:
                        mapped_output = {}
                        for target_key, source_key in self.output_mapping.items():
                            if source_key in output:
                                mapped_output[target_key] = output[source_key]
                        return mapped_output
                    return output
                elif status["status"] == "failed":
                    raise RuntimeError(f"Agent '{self.agent_name}' failed: {status.get('error_message')}")
                return {}
            await asyncio.sleep(1)

        raise TimeoutError(f"Agent '{self.agent_name}' timed out after {self.timeout_seconds}s")
