"""
A2A Server — 基于 DeepSeek 的问答 Agent
对外暴露 A2A 协议接口，让其他 Agent 可以发现自己并提交任务。
"""
import os
import uvicorn
from starlette.applications import Starlette

from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.routes import create_agent_card_routes, create_jsonrpc_routes
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.server.tasks import InMemoryTaskStore, TaskUpdater
from a2a.types import (
    AgentCapabilities, AgentCard, AgentInterface, AgentSkill,
    TaskState,
)
from a2a.helpers import (
    get_message_text, new_task_from_user_message,
    new_text_message, new_text_part,
)
from openai import AsyncOpenAI

# ── DeepSeek Agent ─────────────────────────────────────────

DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY", "你的APIKEY")


class MyAgent:
    """用 DeepSeek 回答问题的 Agent。"""

    def __init__(self):
        self.client = AsyncOpenAI(
            api_key=DEEPSEEK_API_KEY,
            base_url="https://api.deepseek.com/v1",
        )

    async def invoke(self, user_request: str) -> str:
        response = await self.client.chat.completions.create(
            model="deepseek-chat",
            messages=[
                {"role": "system", "content": "你是一个有用的 AI 助手，请用简洁的中文回答用户的问题。"},
                {"role": "user", "content": user_request},
            ],
            temperature=0.7,
            max_tokens=1024,
        )
        return response.choices[0].message.content


# ── Agent Executor ─────────────────────────────────────────

class MyAgentExecutor(AgentExecutor):
    """A2A Executor：把 A2A 任务转发给 MyAgent 处理。"""

    def __init__(self):
        self.agent = MyAgent()

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        # 1. 获取或创建 Task
        if context.current_task:
            task = context.current_task
        else:
            task = new_task_from_user_message(context.message)
            await event_queue.enqueue_event(task)

        # 2. 更新状态：正在处理
        updater = TaskUpdater(event_queue=event_queue, task_id=task.id, context_id=task.context_id)
        await updater.start_work(message=new_text_message("正在处理..."))

        # 3. 提取用户输入，调用 DeepSeek
        query = get_message_text(context.message)
        result = await self.agent.invoke(user_request=query) if query else "没有收到有效输入"

        # 4. 把结果作为 Artifact 返回
        await updater.add_artifact(parts=[new_text_part(text=result, media_type="text/plain")])

        # 5. 标记任务完成
        await updater.complete(message=new_text_message("处理完成！"))
        print(f"  ✓ 处理完成: {result[:60]}...")

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        raise NotImplementedError("暂不支持取消")


# ── Agent Card ─────────────────────────────────────────────

skill = AgentSkill(
    id="qa_bot",
    name="问答机器人",
    description="一个简单的问答 Agent，接收文本问题并返回 DeepSeek 的回答",
    input_modes=["text/plain"],
    output_modes=["text/plain"],
    tags=["问答", "demo", "deepseek"],
    examples=["你好", "什么是 A2A 协议？"],
)

agent_card = AgentCard(
    name="问答 Agent",
    description="一个通用问答 Agent，基于 DeepSeek 回答各种问题",
    version="1.0.0",
    default_input_modes=["text/plain"],
    default_output_modes=["text/plain"],
    capabilities=AgentCapabilities(streaming=True),
    supported_interfaces=[
        AgentInterface(
            protocol_binding="JSONRPC",
            url="http://127.0.0.1:9999",
            protocol_version="1.0",
        ),
    ],
    skills=[skill],
)


# ── 启动服务 ───────────────────────────────────────────────

if __name__ == "__main__":
    request_handler = DefaultRequestHandler(
        agent_executor=MyAgentExecutor(),
        task_store=InMemoryTaskStore(),
        agent_card=agent_card,
    )

    routes = []
    routes.extend(create_agent_card_routes(agent_card))
    routes.extend(create_jsonrpc_routes(request_handler, "/"))

    app = Starlette(routes=routes)
    print("🚀 A2A Server 启动 → http://127.0.0.1:9999")
    print(f"   Agent Card: http://127.0.0.1:9999/.well-known/agent-card.json")
    uvicorn.run(app, host="127.0.0.1", port=9999)