"""
翻译 Agent — 将中文翻译为英文
作为独立 A2A Server 运行，暴露翻译技能
"""
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
from a2a.helpers import (
    get_message_text, new_task_from_user_message,
    new_text_message, new_text_part,
)
from openai import AsyncOpenAI

DEEPSEEK_API_KEY = "你的APIKEY"


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

    async def translate(self, text: str) -> str:
        response = await self.client.chat.completions.create(
            model="deepseek-chat",
            messages=[
                {"role": "system", "content": "你是一个专业的中译英翻译。只输出翻译结果，不要添加解释。"},
                {"role": "user", "content": f"把以下中文翻译成英文：{text}"},
            ],
            temperature=0.3,
            max_tokens=1024,
        )
        return response.choices[0].message.content


class TranslatorExecutor(AgentExecutor):
    def __init__(self):
        self.agent = TranslatorAgent()

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        task = context.current_task or new_task_from_user_message(context.message)
        if not context.current_task:
            await event_queue.enqueue_event(task)

        updater = TaskUpdater(event_queue=event_queue, task_id=task.id, context_id=task.context_id)
        await updater.start_work(message=new_text_message("正在翻译..."))

        query = get_message_text(context.message)
        result = await self.agent.translate(query) if query else "没有收到输入"
        await updater.add_artifact(parts=[new_text_part(text=result, media_type="text/plain")])
        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="translator",
    name="中译英翻译",
    description="将中文文本翻译成英文",
    input_modes=["text/plain"],
    output_modes=["text/plain"],
    tags=["翻译", "demo"],
    examples=["你好吗？", "今天天气真好"],
)

agent_card = AgentCard(
    name="翻译 Agent",
    description="一个中译英翻译 Agent，接收中文返回英文",
    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:9998", protocol_version="1.0"),
    ],
    skills=[skill],
)

if __name__ == "__main__":
    handler = DefaultRequestHandler(
        agent_executor=TranslatorExecutor(),
        task_store=InMemoryTaskStore(),
        agent_card=agent_card,
    )
    routes = []
    routes.extend(create_agent_card_routes(agent_card))
    routes.extend(create_jsonrpc_routes(handler, "/"))
    app = Starlette(routes=routes)
    print("🌐 翻译 Agent → http://127.0.0.1:9998")
    uvicorn.run(app, host="127.0.0.1", port=9998)