"""Persistent OpenClaw gateway WebSocket client with Ed25519 device identity auth.

Maintains a long-lived connection to the local OpenClaw gateway, reused across
all chat requests. Supports streaming via async callbacks and automatic reconnection.
"""

import asyncio
import json
import logging
import uuid
from typing import Callable, Awaitable
from urllib.parse import urlparse

from websockets import connect as ws_connect
from websockets.exceptions import WebSocketException

from bamboo_claw_client.device_identity import (
    DeviceIdentity,
    build_and_sign_device_payload,
    clear_device_token,
    load_device_token,
    load_or_create_identity,
    store_device_token,
)

logger = logging.getLogger(__name__)

CONNECT_TIMEOUT = 10
MAX_BACKOFF = 60


class _ConnectRejected(Exception):
    pass


class GatewayClient:
    """Persistent WebSocket client for the local OpenClaw gateway."""

    PROTOCOL_VERSION = 3
    CLIENT_ID = "openclaw-control-ui"
    CLIENT_VERSION = "1.0.0"
    CLIENT_MODE = "webchat"
    ROLE = "operator"
    SCOPES = [
        "operator.admin",
        "operator.approvals",
        "operator.pairing",
        "operator.write",
    ]

    def __init__(self, url: str, token: str, data_dir: str = "~/.bamboo_claw"):
        self.url = url
        self.token = token
        self.data_dir = data_dir
        self.ws = None
        self.connected = False
        self._identity: DeviceIdentity | None = None
        self._reconnect_task: asyncio.Task | None = None
        self._on_status_change: Callable[[bool], Awaitable] | None = None

    def _get_identity(self) -> DeviceIdentity:
        if self._identity is None:
            self._identity = load_or_create_identity(self.data_dir)
        return self._identity

    async def connect(self) -> dict:
        """Connect to the gateway with device identity challenge-response."""
        identity = self._get_identity()
        stored = load_device_token(self.data_dir, identity.device_id, self.ROLE)
        use_token = stored or self.token
        can_fallback = bool(stored and self.token and stored != self.token)

        try:
            result = await asyncio.wait_for(
                self._attempt_connect(use_token),
                timeout=CONNECT_TIMEOUT,
            )
            self.connected = True
            return result
        except _ConnectRejected:
            if not can_fallback:
                raise
            logger.info("Device token rejected, falling back to shared token")
            clear_device_token(self.data_dir, identity.device_id, self.ROLE)
            result = await asyncio.wait_for(
                self._attempt_connect(self.token),
                timeout=CONNECT_TIMEOUT,
            )
            self.connected = True
            return result

    async def _attempt_connect(self, token: str) -> dict:
        if self.ws:
            try:
                await self.ws.close()
            except Exception:
                pass
            self.ws = None

        try:
            parsed = urlparse(self.url)
            scheme = "https" if parsed.scheme == "wss" else "http"
            origin = f"{scheme}://{parsed.netloc}"

            self.ws = await ws_connect(
                self.url,
                # OpenClaw deployments may sit behind proxies that drop/control
                # websocket ping frames; use app-level reconnect instead.
                ping_interval=None,
                ping_timeout=None,
                additional_headers={"Origin": origin},
            )
        except (OSError, WebSocketException) as e:
            raise Exception(f"Failed to connect to gateway at {self.url}: {e}")

        challenge = await self._receive()
        if challenge.get("event") != "connect.challenge":
            raise Exception(f"Expected connect.challenge, got: {challenge.get('event', 'unknown')}")

        nonce = (challenge.get("payload") or {}).get("nonce", "")
        identity = self._get_identity()
        device = build_and_sign_device_payload(
            identity,
            nonce=nonce,
            token=token,
            client_id=self.CLIENT_ID,
            client_mode=self.CLIENT_MODE,
            role=self.ROLE,
            scopes=self.SCOPES,
        )

        connect_req = {
            "type": "req",
            "id": str(uuid.uuid4()),
            "method": "connect",
            "params": {
                "minProtocol": self.PROTOCOL_VERSION,
                "maxProtocol": self.PROTOCOL_VERSION,
                "client": {
                    "id": self.CLIENT_ID,
                    "version": self.CLIENT_VERSION,
                    "platform": "server",
                    "mode": self.CLIENT_MODE,
                },
                "role": self.ROLE,
                "scopes": self.SCOPES,
                "device": device,
                "caps": [],
                "auth": {"token": token},
            },
        }
        await self.ws.send(json.dumps(connect_req))

        response = await self._receive()
        if not response.get("ok"):
            error = response.get("error", {})
            msg = error.get("message", "Unknown error") if isinstance(error, dict) else str(error)
            await self.ws.close()
            self.ws = None
            raise _ConnectRejected(f"Gateway authentication failed: {msg}")

        auth = (response.get("payload") or {}).get("auth") or {}
        device_token = auth.get("deviceToken")
        if device_token:
            store_device_token(self.data_dir, identity.device_id, self.ROLE, device_token)

        return response

    async def send_chat_message(
        self,
        message: str,
        session_key: str | None = None,
        attachments: list[dict] | None = None,
        on_event: Callable[[dict], Awaitable] | None = None,
    ) -> dict:
        """Send a message and stream events via on_event callback.

        The on_event callback receives dicts like:
          {"state": "delta", "content": "partial..."}
          {"state": "final", "content": "complete text"}
          {"state": "error", "error": "..."}
          {"state": "aborted", "content": "..."}
        """
        if not self.ws or not self.connected:
            raise Exception("Gateway not connected")

        req_id = str(uuid.uuid4())
        run_id = str(uuid.uuid4())
        effective_session_key = session_key or f"agent:main:webchat:dm:web-{uuid.uuid4()}"

        params: dict = {
            "sessionKey": effective_session_key,
            "message": message,
            "deliver": False,
            "idempotencyKey": run_id,
        }

        if attachments:
            api_attachments = []
            extra_urls: list[str] = []
            for att in attachments:
                if isinstance(att, dict) and att.get("type") == "image":
                    api_attachments.append(att)
                elif isinstance(att, dict) and att.get("url"):
                    extra_urls.append(att["url"])
                elif isinstance(att, str):
                    extra_urls.append(att)
            if api_attachments:
                params["attachments"] = api_attachments
            if extra_urls:
                url_text = "\n".join(extra_urls)
                params["message"] = f"{params['message']}\n\n{url_text}" if params["message"] else url_text

        request = {
            "type": "req",
            "id": req_id,
            "method": "chat.send",
            "params": params,
        }
        await self.ws.send(json.dumps(request))

        while True:
            msg = await self._receive()
            if msg.get("type") == "res" and msg.get("id") == req_id:
                if not msg.get("ok"):
                    error = msg.get("error", {})
                    error_msg = error.get("message", "Unknown error") if isinstance(error, dict) else str(error)
                    raise Exception(f"Chat send error: {error_msg}")
                break
            if msg.get("type") == "event":
                continue

        content = ""
        final_session_key = effective_session_key

        while True:
            msg = await self._receive()

            if msg.get("type") == "event" and msg.get("event") == "chat":
                payload = msg.get("payload", {})
                if payload.get("runId") != run_id:
                    continue

                event_session_key = payload.get("sessionKey")
                if event_session_key:
                    final_session_key = event_session_key

                state = payload.get("state")
                message_data = payload.get("message", {})
                text = self._extract_text_content(message_data)

                if state == "delta":
                    content = text
                    if on_event:
                        await on_event({"state": "delta", "content": content})

                elif state == "final":
                    content = text
                    if on_event:
                        await on_event({"state": "final", "content": content})
                    return {"content": content, "sessionKey": final_session_key}

                elif state == "aborted":
                    content = content or "(aborted)"
                    if on_event:
                        await on_event({"state": "aborted", "content": content})
                    return {"content": content, "sessionKey": final_session_key}

                elif state == "error":
                    error_msg = payload.get("errorMessage", "Unknown chat error")
                    if on_event:
                        await on_event({"state": "error", "error": error_msg})
                    raise Exception(f"Chat error: {error_msg}")

    def _extract_text_content(self, message_data: dict) -> str:
        content_list = message_data.get("content", [])
        if isinstance(content_list, list):
            texts = []
            for item in content_list:
                if isinstance(item, dict) and item.get("type") == "text":
                    texts.append(item.get("text", ""))
            return "".join(texts)
        return ""

    async def _receive(self) -> dict:
        raw = await self.ws.recv()
        return json.loads(raw)

    async def disconnect(self):
        self.connected = False
        if self._reconnect_task and not self._reconnect_task.done():
            self._reconnect_task.cancel()
        if self.ws:
            try:
                await self.ws.close()
            except Exception:
                pass
            self.ws = None

    async def reconnect_loop(self):
        """Reconnect with exponential backoff. Runs until successful or cancelled."""
        backoff = 1
        while True:
            logger.info("Attempting gateway reconnect in %ds...", backoff)
            await asyncio.sleep(backoff)
            try:
                await self.connect()
                logger.info("Gateway reconnected successfully")
                if self._on_status_change:
                    await self._on_status_change(True)
                return
            except Exception as e:
                logger.warning("Gateway reconnect failed: %s", e)
                backoff = min(backoff * 2, MAX_BACKOFF)
