# platform/deskclaw/adapters/openclaw_adapter.py
# ═══════════════════════════════════════════════════════════════
# OpenClaw AgentAdapter — 直连 OpenClaw Gateway 的原生适配器
#
# 通过 WebSocket 直连 OpenClaw Gateway，支持：
#   - agent RPC 调用（消息、模型、温度、thinking 等全部参数）
#   - agent 管理（agents.list, agents.files.set 等）
#   - 控制变量实验（agentId、thinking、sessionKey、workspace files）
#
# 协议参考：
#   - Transport: raw WebSocket (text frames, JSON)
#   - Default port: 18789
#   - Protocol version: 3
#   - Handshake: connect.challenge event → connect request → hello-ok
# ═══════════════════════════════════════════════════════════════

from __future__ import annotations

import asyncio
import json
import logging
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, AsyncIterator

import websockets
import websockets.asyncio.client

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

logger = logging.getLogger("deskclaw.adapters.openclaw")

_CLIENT_VERSION = "0.2.2"
_PROTOCOL_VERSION = 3


# ── Configuration ──


@dataclass
class OpenClawConfig:
    """Configuration for an OpenClaw adapter instance."""

    target: str = "ws://localhost:18789"
    token: str = ""
    agent_id: str = ""
    thinking: str = ""
    session_key: str = ""
    scopes: list[str] = field(
        default_factory=lambda: ["operator.read", "operator.write", "operator.admin"]
    )
    request_timeout_s: float = 300.0
    handshake_timeout_s: float = 10.0
    context_window: int = 128_000
    cost_per_1k_input: float = 0.0
    cost_per_1k_output: float = 0.0


# ── Adapter ──


@register_adapter("openclaw")
class OpenClawAdapter(AgentAdapter):
    """
    直连 OpenClaw Gateway 的 AgentAdapter 实现。

    维护一个持久 WebSocket 连接，通过 OpenClaw 原生 RPC 协议
    驱动 agent 执行。支持 set_run_context() 注入控制变量实验
    的 per-run 参数（agentId、thinking、sessionKey 等）。
    """

    def __init__(
        self,
        target: str = "ws://localhost:18789",
        token: str = "",
        agent_id: str = "",
        thinking: str = "",
        **_kwargs: Any,
    ) -> None:
        self._config = OpenClawConfig(
            target=target,
            token=token,
            agent_id=agent_id,
            thinking=thinking,
        )
        self._ws: websockets.asyncio.client.ClientConnection | None = None
        self._pending: dict[str, asyncio.Future[dict]] = {}
        self._reader_task: asyncio.Task[None] | None = None
        self._connected: asyncio.Event | None = None

        self._run_context: dict[str, Any] = {}

    @classmethod
    def from_config(cls, config: OpenClawConfig) -> OpenClawAdapter:
        adapter = cls(
            target=config.target,
            token=config.token,
            agent_id=config.agent_id,
            thinking=config.thinking,
        )
        adapter._config = config
        return adapter

    # ── AgentAdapter protocol ──

    def name(self) -> str:
        agent_suffix = self._effective_agent_id()
        return f"openclaw:{agent_suffix}" if agent_suffix else "openclaw"

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

    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:
        await self._ensure_connected()

        workspace_files = self._run_context.get("workspace_files")
        if workspace_files and isinstance(workspace_files, dict):
            await self._apply_workspace_files(workspace_files)

        start = time.monotonic()
        payload = await self._rpc_agent(
            messages, max_tokens=max_tokens, temperature=temperature, timeout_s=timeout_s
        )
        latency_ms = int((time.monotonic() - start) * 1000)

        content = payload.get("content", "")
        usage = payload.get("usage", {})

        assistant_msg = AgentMessage(
            role="assistant",
            content=content,
            metadata={
                "adapter": "openclaw",
                "agent_id": self._effective_agent_id(),
                "model": payload.get("model", ""),
                "thinking_content": payload.get("thinking", ""),
            },
        )

        return AgentResponse(
            messages=list(messages) + [assistant_msg],
            final_answer=content,
            trace={
                "raw_payload": payload,
                "adapter": "openclaw",
                "agent_id": self._effective_agent_id(),
                "thinking_mode": self._effective_thinking(),
                "thinking_content": payload.get("thinking", ""),
                "model": payload.get("model", ""),
                "session_key": self._effective_session_key(),
            },
            usage={
                "input_tokens": usage.get("prompt_tokens", usage.get("input_tokens", 0)),
                "output_tokens": usage.get("completion_tokens", usage.get("output_tokens", 0)),
                "api_calls": 1,
            },
            latency_ms=latency_ms,
        )

    async def stream(
        self,
        messages: list[AgentMessage],
        **kwargs: Any,
    ) -> AsyncIterator[AgentMessage]:
        await self._ensure_connected()
        payload = await self._rpc_agent(messages, **kwargs)
        yield AgentMessage(role="assistant", content=payload.get("content", ""))

    def max_context_window(self) -> int:
        return self._config.context_window

    def cost_per_1k_tokens(self) -> dict[str, float]:
        return {
            "input": self._config.cost_per_1k_input,
            "output": self._config.cost_per_1k_output,
        }

    # ── per-run context (controlled variable experiments) ──

    def set_run_context(self, metadata: dict[str, Any]) -> None:
        """
        注入当前 run 的 OpenClaw 特有参数。
        由 TemplateAdapter 在 execute() 前调用。

        支持的 key：
          - agent_id:        目标 agentId
          - thinking:        思考模式 ("", "low", "medium", "high")
          - session_key:     会话隔离 key
          - workspace_files: dict[filename, content] 在 run 前写入 agent workspace
        """
        self._run_context = dict(metadata)
        logger.debug("Run context set: %s", list(self._run_context.keys()))

    def clear_run_context(self) -> None:
        self._run_context = {}

    # ── connection lifecycle ──

    async def connect(self) -> None:
        if self._ws and self._connected and self._connected.is_set():
            return
        self._connected = asyncio.Event()
        self._ws = await websockets.asyncio.client.connect(self._config.target)
        self._reader_task = asyncio.create_task(self._reader_loop())
        try:
            await asyncio.wait_for(
                self._connected.wait(), timeout=self._config.handshake_timeout_s
            )
        except asyncio.TimeoutError:
            await self.close()
            raise ConnectionError(
                f"OpenClaw handshake timed out ({self._config.target})"
            )
        logger.info("OpenClaw adapter connected to %s", self._config.target)

    async def close(self) -> None:
        if self._reader_task and not self._reader_task.done():
            self._reader_task.cancel()
            try:
                await self._reader_task
            except asyncio.CancelledError:
                pass
        if self._ws:
            await self._ws.close()
            self._ws = None
        for fut in self._pending.values():
            if not fut.done():
                fut.cancel()
        self._pending.clear()
        self._run_context = {}
        logger.info("OpenClaw adapter closed")

    @property
    def is_connected(self) -> bool:
        return (
            self._ws is not None
            and self._connected is not None
            and self._connected.is_set()
        )

    # ── agent management RPC ──

    async def list_agents(self) -> list[dict[str, Any]]:
        await self._ensure_connected()
        res = await self._rpc("agents.list", {})
        return res.get("agents", [])

    async def get_agent_files(self, agent_id: str) -> list[dict[str, Any]]:
        await self._ensure_connected()
        res = await self._rpc("agents.files.list", {"agentId": agent_id})
        return res.get("files", [])

    async def get_agent_file(self, agent_id: str, filename: str) -> str:
        await self._ensure_connected()
        res = await self._rpc("agents.files.get", {"agentId": agent_id, "name": filename})
        return res.get("file", {}).get("content", "")

    async def set_agent_file(self, agent_id: str, filename: str, content: str) -> None:
        await self._ensure_connected()
        await self._rpc(
            "agents.files.set",
            {"agentId": agent_id, "name": filename, "content": content},
        )
        logger.debug("Set agent file %s/%s (%d chars)", agent_id, filename, len(content))

    # ── internals ──

    def _effective_agent_id(self) -> str:
        return self._run_context.get("agent_id", "") or self._config.agent_id

    def _effective_thinking(self) -> str:
        return self._run_context.get("thinking", "") or self._config.thinking

    def _effective_session_key(self) -> str:
        ctx_key = self._run_context.get("session_key", "")
        if ctx_key:
            return ctx_key
        if self._config.session_key:
            return self._config.session_key
        return f"deskclaw-exp-{uuid.uuid4().hex[:8]}"

    async def _ensure_connected(self) -> None:
        if not self.is_connected:
            await self.connect()

    async def _rpc_agent(
        self,
        messages: list[AgentMessage],
        *,
        max_tokens: int = 4096,
        temperature: float = 0.0,
        timeout_s: int = 300,
        **_kwargs: Any,
    ) -> dict[str, Any]:
        """Send an `agent` RPC call, wait for completion, fetch chat history."""
        user_message = ""
        system_parts: list[str] = []
        for m in messages:
            if m.role == "user":
                user_message = m.content
            elif m.role == "system":
                system_parts.append(m.content)

        if not user_message:
            user_message = messages[-1].content if messages else ""

        session_key = self._effective_session_key()

        params: dict[str, Any] = {
            "message": user_message,
            "idempotencyKey": uuid.uuid4().hex,
            "sessionKey": session_key,
        }

        agent_id = self._effective_agent_id()
        if agent_id:
            params["agentId"] = agent_id

        thinking = self._effective_thinking()
        if thinking:
            params["thinking"] = thinking

        if system_parts:
            params["extraSystemPrompt"] = "\n\n".join(system_parts)

        params["timeout"] = timeout_s

        res = await self._rpc("agent", params, timeout=timeout_s + 30)

        run_id = res.get("runId", "")
        if run_id:
            await self._wait_for_agent_run(run_id, timeout_s)

        return await self._fetch_assistant_response(session_key, res)

    async def _wait_for_agent_run(
        self, run_id: str, timeout_s: int
    ) -> dict[str, Any]:
        """Wait for an agent run to complete."""
        try:
            return await self._rpc(
                "agent.wait",
                {"runId": run_id},
                timeout=timeout_s + 30,
            )
        except TimeoutError:
            logger.warning("agent.wait timed out for runId=%s", run_id)
            return {"status": "timeout"}

    async def _fetch_assistant_response(
        self, session_key: str, run_info: dict[str, Any]
    ) -> dict[str, Any]:
        """Fetch chat history and extract the last assistant response."""
        try:
            history = await self._rpc(
                "chat.history",
                {"sessionKey": session_key},
                timeout=30,
            )
        except Exception as exc:
            logger.warning("Failed to fetch chat.history: %s", exc)
            return {"content": "", "usage": {}, "run_info": run_info}

        assistant_msgs = [
            m for m in history.get("messages", []) if m.get("role") == "assistant"
        ]
        if not assistant_msgs:
            return {"content": "", "usage": {}, "run_info": run_info}

        last = assistant_msgs[-1]
        text_parts: list[str] = []
        thinking_parts: list[str] = []
        for block in last.get("content", []):
            if isinstance(block, dict):
                if block.get("type") == "text":
                    text_parts.append(block.get("text", ""))
                elif block.get("type") == "thinking":
                    thinking_parts.append(block.get("thinking", ""))
            elif isinstance(block, str):
                text_parts.append(block)

        return {
            "content": "\n".join(text_parts),
            "thinking": "\n".join(thinking_parts),
            "model": last.get("model", ""),
            "provider": last.get("provider", ""),
            "stopReason": last.get("stopReason", ""),
            "usage": history.get("usage", {}),
            "run_info": run_info,
        }

    async def _rpc(
        self,
        method: str,
        params: dict[str, Any],
        *,
        timeout: float | None = None,
    ) -> dict[str, Any]:
        """Send a generic RPC request and await the response."""
        assert self._ws is not None, "Not connected"

        req_id = str(uuid.uuid4())
        frame = {"type": "req", "id": req_id, "method": method, "params": params}

        future: asyncio.Future[dict] = asyncio.get_running_loop().create_future()
        self._pending[req_id] = future

        await self._ws.send(json.dumps(frame))
        logger.debug("RPC → %s (id=%s)", method, req_id)

        effective_timeout = timeout or self._config.request_timeout_s
        try:
            res = await asyncio.wait_for(future, timeout=effective_timeout)
        except asyncio.TimeoutError:
            self._pending.pop(req_id, None)
            raise TimeoutError(f"OpenClaw RPC {method} timed out ({effective_timeout}s)")

        if not res.get("ok", False):
            err = res.get("error", {})
            raise RuntimeError(
                f"OpenClaw RPC error [{method}] "
                f"[{err.get('code', '?')}]: {err.get('message', 'unknown')}"
            )

        return res.get("payload", res)

    async def _apply_workspace_files(self, files: dict[str, str]) -> None:
        """Write workspace files to the target agent before a run."""
        agent_id = self._effective_agent_id()
        if not agent_id:
            logger.warning("workspace_files specified but no agent_id — skipping")
            return
        for filename, content in files.items():
            await self.set_agent_file(agent_id, filename, content)

    # ── WebSocket transport ──

    async def _reader_loop(self) -> None:
        assert self._ws is not None
        try:
            async for raw in self._ws:
                try:
                    frame = json.loads(raw)
                except (json.JSONDecodeError, TypeError):
                    logger.warning("Non-JSON frame received, ignoring")
                    continue
                await self._dispatch(frame)
        except websockets.exceptions.ConnectionClosed:
            logger.warning("OpenClaw WebSocket closed unexpectedly")
            if self._connected:
                self._connected.clear()
        except asyncio.CancelledError:
            return

    async def _dispatch(self, frame: dict) -> None:
        ftype = frame.get("type", "")

        if ftype == "event":
            await self._handle_event(frame)
        elif ftype == "res":
            fid = frame.get("id", "")
            future = self._pending.pop(fid, None)
            if future and not future.done():
                future.set_result(frame)
        else:
            logger.debug("Unknown frame type: %s", ftype)

    async def _handle_event(self, frame: dict) -> None:
        event_name = frame.get("event", "")

        if event_name == "connect.challenge":
            nonce = frame.get("payload", {}).get("nonce", "")
            await self._send_connect(nonce)
        elif event_name == "tick":
            pass
        else:
            logger.debug("Event: %s", event_name)

    async def _send_connect(self, nonce: str) -> None:
        """Send the connect handshake frame.

        The response is handled by ``_dispatch`` which will resolve the
        future registered here.  A callback on that future sets
        ``_connected`` — this avoids deadlocking the reader loop (since
        this method is called *from within* the reader loop coroutine).
        """
        assert self._ws is not None
        req_id = str(uuid.uuid4())
        params: dict[str, Any] = {
            "minProtocol": _PROTOCOL_VERSION,
            "maxProtocol": _PROTOCOL_VERSION,
            "client": {
                "id": "gateway-client",
                "version": _CLIENT_VERSION,
                "platform": "deskclaw-arena",
                "mode": "backend",
            },
            "role": "operator",
            "scopes": self._config.scopes,
        }
        if self._config.token:
            params["auth"] = {"token": self._config.token}

        frame = {"type": "req", "id": req_id, "method": "connect", "params": params}

        future: asyncio.Future[dict] = asyncio.get_running_loop().create_future()
        self._pending[req_id] = future

        def _on_connect_response(fut: asyncio.Future[dict]) -> None:
            if fut.cancelled():
                return
            res = fut.result()
            if res.get("ok"):
                self._connected.set()  # type: ignore[union-attr]
                logger.info("OpenClaw handshake succeeded")
            else:
                err = res.get("error", {})
                logger.error(
                    "OpenClaw connect rejected: %s", err.get("message", "unknown")
                )

        future.add_done_callback(_on_connect_response)

        await self._ws.send(json.dumps(frame))
        logger.debug("Sent connect request (nonce=%s)", nonce)

