from python_a2a import A2AServer, run_server, AgentCard, AgentSkill, TaskStatus, TaskState
from python_a2a.mcp import MCPClient
import asyncio

# A2A Agent 的名片
agent_card = AgentCard(
    name="WeatherServer",
    description="用来查询天气",
    url="http://127.0.0.1:8005",
    skills=[AgentSkill(name="查询天气", description="查询指定城市的天气")]
)

class WeatherServer(A2AServer):
    def __init__(self):
        super().__init__(agent_card=agent_card)
        # 创建一个MCPClient对象
        self.mcp_client = MCPClient('http://127.0.0.1:6005')

    def handle_task(self, task):
        print("收到A2A任务的task:=>", task)
        # 获取任务内容
        query = (task.message or {}).get('content', {}).get('text', '')

        # 具体处理逻辑
        if "天气" in query:
            # 获取城市
            city = "北京"
            # 调用MCP Server中的工具
            # 因为这个工具是异步的，所以需要使用 asyncio.run() 进行调用，否则只会返回一个调用的对象！！
            weather_result = asyncio.run(self.mcp_client.call_tool(tool_name="get_weather", city=city))
            print("天气查询结果:=>", weather_result)
            # 将查询结果放到 task.artifacts
            task.artifacts = [{"parts": [{"type": "text", "text": weather_result}]}]
        else:
            task.artifacts = [{"parts": [{"type": "text", "text": "无法理解的任务"}]}]

        # 设置任务状态为完成
        task.status = TaskStatus(TaskState.COMPLETED)

        print(f"[{self.agent_card.name} 日志] 任务完成，结果已返回给 A2A")
        print("task:=>",task)
        print("task.artifacts:=>",task.artifacts)

        return task


if __name__ == '__main__':
    run_server(WeatherServer(), host="127.0.0.1", port=8005)