"""OpenClaw WebSocket bridge.

Translates OpenAI chat-completion requests into OpenClaw's native WebSocket
protocol (JSON text frames with req/res/event message types).

Protocol reference:
  - Transport: raw WebSocket (not Socket.IO)
  - Default port: 18789
  - Frame types: req, res, event
  - Handshake: connect.challenge event → connect request → hello-ok response
"""

from __future__ import annotations

import asyncio
import json
import logging
import uuid
from typing import Any

import websockets
import websockets.asyncio.client

from deskclaw_client.bridges import Bridge, register_bridge

log = logging.getLogger(__name__)


@register_bridge("openclaw")
class OpenClawBridge(Bridge):
    """Bridge that speaks OpenClaw's WebSocket protocol."""

    def __init__(
        self,
        target: str = "ws://localhost:18789",
        token: str | None = None,
        **_kwargs: Any,
    ) -> None:
        self._target = target
        self._token = token
        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()

    # ── lifecycle ──

    async def connect(self) -> None:
        self._ws = await websockets.asyncio.client.connect(self._target)
        self._reader_task = asyncio.create_task(self._reader_loop())
        try:
            await asyncio.wait_for(self._connected.wait(), timeout=10.0)
        except asyncio.TimeoutError:
            await self.close()
            raise ConnectionError(
                f"OpenClaw handshake timed out ({self._target})"
            )
        log.info("OpenClaw bridge connected to %s", self._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()
        log.info("OpenClaw bridge closed")

    # ── chat ──

    async def chat(
        self,
        messages: list[dict[str, Any]],
        *,
        model: str = "",
        temperature: float = 1.0,
        max_tokens: int | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        if not self._ws:
            raise RuntimeError("Bridge not connected — call connect() first")

        req_id = str(uuid.uuid4())
        params: dict[str, Any] = {"messages": messages}
        if model:
            params["model"] = model
        if max_tokens is not None:
            params["max_tokens"] = max_tokens
        params["temperature"] = temperature

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

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

        await self._ws.send(json.dumps(frame))
        log.debug("Sent req %s", req_id)

        try:
            res = await asyncio.wait_for(future, timeout=300.0)
        except asyncio.TimeoutError:
            self._pending.pop(req_id, None)
            raise TimeoutError(f"OpenClaw request {req_id} timed out (300s)")

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

        return self._to_openai_response(res.get("payload", {}), model)

    async def health(self) -> dict[str, Any]:
        return {
            "status": "ok" if self._ws and self._connected.is_set() else "disconnected",
            "target": self._target,
            "bridge": "openclaw",
        }

    # ── internals ──

    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):
                    log.warning("Non-JSON frame received, ignoring")
                    continue
                await self._dispatch(frame)
        except websockets.exceptions.ConnectionClosed:
            log.warning("OpenClaw WebSocket closed unexpectedly")
        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:
                log.debug("Unmatched response id=%s", fid)
        else:
            log.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:
            log.debug("Event: %s", event_name)

    async def _send_connect(self, nonce: str) -> None:
        assert self._ws is not None
        req_id = str(uuid.uuid4())
        params: dict[str, Any] = {
            "minProtocol": 3,
            "maxProtocol": 3,
            "client": {
                "id": "deskclaw-client",
                "version": "0.2.0",
                "platform": "bridge",
                "mode": "operator",
            },
            "role": "operator",
            "scopes": ["operator.read", "operator.write"],
            "device": {
                "id": f"deskclaw-{uuid.uuid4().hex[:8]}",
                "nonce": nonce,
            },
        }
        if self._token:
            params["auth"] = {"token": self._token}

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

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

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

        try:
            res = await asyncio.wait_for(future, timeout=10.0)
        except asyncio.TimeoutError:
            self._pending.pop(req_id, None)
            raise ConnectionError("OpenClaw connect handshake timed out")

        if res.get("ok"):
            self._connected.set()
            log.info("OpenClaw handshake succeeded")
        else:
            err = res.get("error", {})
            raise ConnectionError(
                f"OpenClaw connect rejected: {err.get('message', 'unknown')}"
            )

    @staticmethod
    def _to_openai_response(payload: dict, model: str) -> dict[str, Any]:
        """Normalise an OpenClaw payload into the OpenAI chat-completion shape."""
        content = payload.get("content") or payload.get("text") or payload.get("answer", "")
        if isinstance(content, list):
            content = "\n".join(str(c) for c in content)

        return {
            "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
            "object": "chat.completion",
            "model": model or "openclaw",
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": str(content)},
                    "finish_reason": "stop",
                }
            ],
            "usage": payload.get("usage", {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}),
        }
