"""
A2A 客户端测试脚本
通过 A2A 协议调用本地的 Dev Agent 服务。

使用方式：
    python a2a_client.py "你好，介绍一下你自己"
    python a2a_client.py --context my_session "你好"   # 指定 context_id（会话）
"""
import asyncio
import sys
import argparse

from a2a.client import create_client
from a2a.types.a2a_pb2 import SendMessageRequest, Message, Part, Role


A2A_URL = "http://127.0.0.1:8001"


async def main():
    parser = argparse.ArgumentParser(description="A2A 客户端测试")
    parser.add_argument("message", help="发送给 agent 的消息")
    parser.add_argument("--context", "-c", default="", help="context_id，用于会话续传")
    args = parser.parse_args()

    print(f"连接到: {A2A_URL}")
    if args.context:
        print(f"会话: {args.context}")
    print(f"发送: {args.message}")
    print("-" * 40)

    client = await create_client(A2A_URL)
    try:
        # 构建请求
        request = SendMessageRequest(
            message=Message(
                role=Role.ROLE_USER,
                parts=[Part(text=args.message)],
            ),
        )
        if args.context:
            request.message.context_id = args.context

        # 发送消息（流式返回）
        full_text = ""
        context_id = ""
        async for response in client.send_message(request):
            # 流式响应里可能有 message、status_update 等
            if response.HasField("message"):
                msg = response.message
                if msg.context_id:
                    context_id = msg.context_id
                for part in msg.parts:
                    if part.text:
                        full_text += part.text

        print("Agent 回复:")
        print(full_text)
        if context_id:
            print(f"\n(context_id: {context_id})")
    finally:
        await client.close()


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