"""
A2A vs 直连 对比测试
交互式输入问题，同时展示两种调用方式的结果

运行前提：问答 Agent (server.py) 必须在 9999 端口运行
"""
import asyncio
import httpx
from datetime import datetime
from openai import AsyncOpenAI
from a2a.client import A2ACardResolver, ClientConfig, create_client
from a2a.helpers import new_text_message
from a2a.types import Role, SendMessageRequest

DEEPSEEK_API_KEY = "你的APIKEY"
A2A_SERVER_URL = "http://127.0.0.1:9999"


def extract_text(task) -> str:
    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 direct_call(question: str) -> str:
    """方式 A：直连 DeepSeek，不经过 A2A"""
    client = AsyncOpenAI(
        api_key=DEEPSEEK_API_KEY,
        base_url="https://api.deepseek.com/v1",
    )
    response = await client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": "你是一个有用的 AI 助手，请用简洁的中文回答用户的问题。"},
            {"role": "user", "content": question},
        ],
        temperature=0.7,
        max_tokens=1024,
    )
    result = response.choices[0].message.content
    await client.close()
    return result


async def a2a_call(question: str) -> str:
    """方式 B：走 A2A 协议 → Server → DeepSeek"""
    async with httpx.AsyncClient() as http_client:
        resolver = A2ACardResolver(httpx_client=http_client, base_url=A2A_SERVER_URL)
        agent_card = await resolver.get_agent_card()

        config = ClientConfig(streaming=False)
        client = await create_client(agent=agent_card, client_config=config)
        message = new_text_message(question, role=Role.ROLE_USER)
        request = SendMessageRequest(message=message)

        result = ""
        async for response in client.send_message(request):
            result = extract_text(response.task)

        await client.close()
        return result


async def main():
    print("=" * 72)
    print("  A2A vs 直连 — 交互式对比测试")
    print("  =============================")
    print("  方式 A：Client → DeepSeek API（直连，无 A2A）")
    print("  方式 B：Client → A2A Server → DeepSeek API（经过 A2A 协议）")
    print("=" * 72)

    # 先检查 A2A Server 是否在运行
    print("\n[*] 检查 A2A Server 状态...", end=" ")
    try:
        async with httpx.AsyncClient() as c:
            r = await c.get(f"{A2A_SERVER_URL}/.well-known/agent-card.json", timeout=3)
            if r.status_code == 200:
                print("✅ 运行中")
            else:
                print(f"⚠️ 响应异常 ({r.status_code})")
    except Exception:
        print("❌ 未运行！请先启动：python server.py")
        return

    while True:
        question = input("\n" + "─" * 72 + "\n输入问题（输入 exit 退出）> ").strip()
        if not question:
            continue
        if question.lower() in ("exit", "quit", "q"):
            print("bye")
            break

        print()

        # 方式 A：直连
        t0 = datetime.now()
        print("  ┌─ [A] 直连 DeepSeek ──────────────────────────────")
        try:
            direct_result = await direct_call(question)
            t1 = datetime.now()
            for line in direct_result.split("\n"):
                print(f"  │ {line}")
            print(f"  └─ ⏱ {(t1 - t0).total_seconds():.2f}s")
        except Exception as e:
            print(f"  │ ❌ {e}")
            print("  └─")

        print()

        # 方式 B：走 A2A
        t2 = datetime.now()
        print("  ┌─ [B] 通过 A2A 协议 ──────────────────────────────")
        print("  │   Client → A2A Server → DeepSeek")
        try:
            a2a_result = await a2a_call(question)
            t3 = datetime.now()
            for line in a2a_result.split("\n"):
                print(f"  │ {line}")
            print(f"  └─ ⏱ {(t3 - t2).total_seconds():.2f}s")
        except Exception as e:
            print(f"  │ ❌ {e}")
            print("  └─")


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