import asyncio
from python_a2a import AgentNetwork, A2AClient, Task, Message, MessageRole, TextContent
import json
import uuid
from time import sleep


async def main():
    # 步骤1：初始化 AgentNetwork 并注册专家 Agent
    # 1.1 创建 AgentNetwork 实例，用于管理 Agent 集合
    network = AgentNetwork(name="TravelOrchestrator")
    # 1.2 添加票务 Agent
    network.add("TicketAgent", "http://127.0.0.1:5009")
    # 1.3 添加天气 Agent
    network.add("WeatherAgent", "http://127.0.0.1:5008")
    print("[主控日志] AgentNetwork 初始化完成，已添加专家代理：")
    for agent_info in network.list_agents():
        print(json.dumps(agent_info, indent=4, ensure_ascii=False))
    print("-" * 50)

    # 步骤2：执行串行任务流
    # 2.1 任务一：查询天气
    weather_query = "北京的天气怎么样"  # 用户请求

    #  获取 WeatherAgent 的客户端
    weather_client = network.get_agent("WeatherAgent")

    # 调用客户端
    # 方式一：直接调用
    weather_result = weather_client.ask(weather_query)
    print("[主控日志] 天气查询结果：")
    print(weather_result)

    # 方式二：异步调用
    # Message用来存储具体的任务内容
    message = Message(content=TextContent(text=weather_query), role=MessageRole.USER)
    # Task中存储和封装Message
    weather_task = Task(message=message.to_dict(), id="task-" + str(uuid.uuid4()))
    # 使用 send_task_async 实现异步调用
    weather_result =  await weather_client.send_task_async(weather_task)

    # 解析结果
    weather_info = "未知天气"
    try:
        # 获取 artifacts 中的文本部分
        weather_parts = weather_result.artifacts[0]["parts"]
        if weather_parts and weather_parts[0].get("type") == "text":
            weather_info = weather_parts[0].get("text")
            print(f"[主控日志] 收到 WeatherAgent 的结果: '{weather_info}'")
    except Exception as e:
        print(f"[主控日志] 解析天气结果出错: {e}")


    # 2.2 任务二： 根据天气的结果预定火车票
    ticket_query = f"预订一张从北京到上海的火车票，当前天气是：{weather_info}"

    #  获取 TicketAgent 的客户端
    ticket_client = network.get_agent("TicketAgent")

    # 异步调用
    # Message用来存储具体的任务内容
    ticket_message = Message(content=TextContent(text=ticket_query), role=MessageRole.USER)
    # Task中存储和封装Message
    ticket_task = Task(message=ticket_message.to_dict(), id="task-" + str(uuid.uuid4()))
    # 使用 send_task_async 实现异步调用
    ticket_result =  await ticket_client.send_task_async(ticket_task)

    # 打印最终结果
    print(f"\n[主控日志] 收到 TicketAgent 的最终结果：")
    print(json.dumps(ticket_result.to_dict(), indent=4, ensure_ascii=False))


if __name__ == "__main__":
    asyncio.run(main())