"""
A2A Demo Client — 展示 A2A 协议核心价值

流程：
1. 发现两个 Agent 的能力（Agent Card）
2. 分别提交任务给不同的 Agent
3. 展示跨 Agent 协作：问答 Agent 回答 → 翻译 Agent 翻译
"""
import asyncio
import httpx
from a2a.client import A2ACardResolver, ClientConfig, create_client
from a2a.helpers import new_text_message, display_agent_card
from a2a.types import Role, SendMessageRequest

AGENTS = {
    "问答": "http://127.0.0.1:9999",
    "翻译": "http://127.0.0.1:9998",
}


def extract_text(task) -> str:
    """从 Task 的 artifacts 中提取回复文本"""
    texts = []
    if task.artifacts:
        for artifact in task.artifacts:
            for part in artifact.parts:
                t = part.text.strip() if part.text else ""
                if t:
                    texts.append(t)
    return "\n".join(texts)


async def main():
    async with httpx.AsyncClient() as http_client:
        # ═══════════════════════════════════════════════
        # Phase 1: 发现多个 Agent 的能力
        # ═══════════════════════════════════════════════
        print("=" * 70)
        print("  Phase 1: 🔍 发现 Agent 能力（Agent Card）")
        print("  A2A 协议让 Client 可以在运行期发现 Agent 能干什么")
        print("=" * 70)

        cards = {}
        for name, url in AGENTS.items():
            print(f"\n  ▶ 发现 {name} Agent → {url}")
            resolver = A2ACardResolver(httpx_client=http_client, base_url=url)
            card = await resolver.get_agent_card()
            cards[name] = card
            # 只取关键信息展示，不打印全量
            skills_str = ", ".join(s.name for s in card.skills) if card.skills else "(无)"
            print(f"     Agent: {card.name}  |  技能: {skills_str}")

        # ═══════════════════════════════════════════════
        # Phase 2: 分别调用不同的 Agent
        # ═══════════════════════════════════════════════
        print("\n" + "=" * 70)
        print("  Phase 2: 🤝 通过相同协议调用不同 Agent")
        print("  不管是问答还是翻译，Client 都走同一套 A2A 接口")
        print("=" * 70)

        for name, card in cards.items():
            print(f"\n  ── {name} Agent ──")
            config = ClientConfig(streaming=False)
            client = await create_client(agent=card, client_config=config)

            # 给每个 Agent 一个适合它的任务
            if name == "问答":
                query = "Explain what A2A protocol is in one sentence"
            else:
                query = "请把下面这句话翻译成英文：人工智能正在改变世界"

            print(f"  输入: {query}")

            message = new_text_message(query, role=Role.ROLE_USER)
            request = SendMessageRequest(message=message)

            async for response in client.send_message(request):
                text = extract_text(response.task)
                if text:
                    # 缩进展示，区分输入输出
                    for line in text.split("\n"):
                        print(f"  输出: {line}")

            await client.close()

        # ═══════════════════════════════════════════════
        # Phase 3: 跨 Agent 协作（这才是 A2A 的精髓）
        # ═══════════════════════════════════════════════
        print("\n" + "=" * 70)
        print("  Phase 3: 🔄 跨 Agent 协作管道")
        print("  问答 Agent 的回答 → 翻译 Agent 翻译成英文")
        print("  不同 Agent 通过标准协议接力，无需定制适配")
        print("=" * 70)

        # 3a. 先问问题
        question = "什么是 A2A 协议？用一句话说"
        print(f"\n  Step 1: 用户提问 → 问答 Agent")
        print(f"  输入: {question}")

        qa_card = cards["问答"]
        qa_client = await create_client(agent=qa_card, client_config=ClientConfig(streaming=False))
        qa_msg = new_text_message(question, role=Role.ROLE_USER)
        async for resp in qa_client.send_message(SendMessageRequest(message=qa_msg)):
            qa_answer = extract_text(resp.task)
        await qa_client.close()
        print(f"  问答 Agent 回复: {qa_answer[:100]}...")

        # 3b. 把回答送去翻译
        print(f"\n  Step 2: 问答结果 → 翻译 Agent")
        print(f"  输入: (上一步的回答，翻译成英文)")

        trans_card = cards["翻译"]
        trans_client = await create_client(agent=trans_card, client_config=ClientConfig(streaming=False))
        trans_msg = new_text_message(f"把下面这句话翻译成英文：{qa_answer}", role=Role.ROLE_USER)
        async for resp in trans_client.send_message(SendMessageRequest(message=trans_msg)):
            trans_result = extract_text(resp.task)
        await trans_client.close()
        print(f"  翻译 Agent 输出: {trans_result[:200]}...")

        # ═══════════════════════════════════════════════
        print("\n" + "=" * 70)
        print("  ✅ A2A Demo 完成")
        print("  🔑 关键点：")
        print("  • 两个 Agent 能力完全不同（问答 vs 翻译）")
        print("  • Client 通过 Agent Card 动态发现它们的能力")
        print("  • Client 用同一套 A2A 协议与两者通信，零适配代码")
        print("  • 跨 Agent 管道自动流转，无需中间格式转换")
        print("=" * 70)


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