# platform/deskclaw/adapters/nanobot.py
# ═══════════════════════════════════════════════════════════════
# NanoBotAdapter — 轻量级 Agent 的 U 盘式接入示范
# 实现此文件 + @register_adapter 即可参与全部测评
# 零修改 DeskClaw 核心代码
# ═══════════════════════════════════════════════════════════════

import json
import time

import httpx

from deskclaw.adapters.protocol import (
    AgentAdapter,
    AgentMessage,
    AgentResponse,
    AgentCapability,
)
from deskclaw.adapters.registry import register_adapter


@register_adapter("nanobot-v1")
class NanoBotAdapter(AgentAdapter):
    """
    NanoBot Agent Adapter.
    演示 U 盘式接入的最简实现：
    - 实现 4 个必要方法
    - 用装饰器注册
    - 无需修改任何 DeskClaw 核心代码
    """

    def __init__(self, api_url: str = "http://localhost:9000", api_key: str = ""):
        self.api_url = api_url.rstrip("/")
        self.api_key = api_key
        self._client = httpx.AsyncClient(
            base_url=self.api_url,
            headers={"Authorization": f"Bearer {self.api_key}"},
            timeout=httpx.Timeout(300.0),
        )

    def name(self) -> str:
        return "nanobot-v1"

    def capabilities(self) -> set[AgentCapability]:
        return {AgentCapability.TEXT_GENERATION, AgentCapability.CODE_EXECUTION}

    async def execute(
        self,
        messages: list[AgentMessage],
        *,
        max_tokens: int = 4096,
        temperature: float = 0.0,
        tools: list[dict] | None = None,
        timeout_s: int = 300,
    ) -> AgentResponse:
        start = time.monotonic()

        payload = {
            "messages": [{"role": m.role, "content": m.content} for m in messages],
            "max_tokens": max_tokens,
            "temperature": temperature,
        }
        if tools:
            payload["tools"] = tools

        resp = await self._client.post("/v1/chat/completions", json=payload)
        resp.raise_for_status()
        data = resp.json()

        latency = int((time.monotonic() - start) * 1000)
        assistant_msg = AgentMessage(
            role="assistant",
            content=data["choices"][0]["message"]["content"],
            tool_calls=data["choices"][0]["message"].get("tool_calls"),
            metadata={"model": data.get("model", "nanobot-v1")},
        )

        return AgentResponse(
            messages=list(messages) + [assistant_msg],
            final_answer=assistant_msg.content,
            trace={"raw_response": data},
            usage={
                "input_tokens": data.get("usage", {}).get("prompt_tokens", 0),
                "output_tokens": data.get("usage", {}).get("completion_tokens", 0),
                "api_calls": 1,
            },
            latency_ms=latency,
        )

    async def stream(self, messages: list[AgentMessage], **kwargs):
        async with self._client.stream(
            "POST",
            "/v1/chat/completions",
            json={
                "messages": [{"role": m.role, "content": m.content} for m in messages],
                "stream": True,
                **kwargs,
            },
        ) as resp:
            async for line in resp.aiter_lines():
                if line.startswith("data: ") and line != "data: [DONE]":
                    chunk = json.loads(line[6:])
                    delta = chunk["choices"][0].get("delta", {})
                    if "content" in delta:
                        yield AgentMessage(role="assistant", content=delta["content"])

    def max_context_window(self) -> int:
        return 32_000  # NanoBot 窗口较小
