"""WebSocket client for connecting to the Bamboo Claw central server.

Handles authentication, heartbeats, command reception, and event relay.
"""

import asyncio
import json
import logging
import platform
from typing import Callable, Awaitable, Optional

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

logger = logging.getLogger(__name__)

HEARTBEAT_INTERVAL = 30
HEARTBEAT_ACK_TIMEOUT = 10
MAX_BACKOFF = 60


class ServerConnection:
    """Persistent WebSocket connection to the Bamboo Claw central server."""

    def __init__(
        self,
        server_url: str,
        token: str,
        client_id: str = "",
        hostname: str = "",
    ):
        self.server_url = server_url
        self.token = token
        self.client_id = client_id
        self.hostname = hostname or platform.node()
        self.ws = None
        self.connected = False
        self._heartbeat_task: asyncio.Task | None = None
        self._on_command: Optional[Callable[[dict], Awaitable]] = None

    async def connect(self) -> str:
        """Connect and authenticate with the server. Returns assigned client_id."""
        from bamboo_claw_client.config import build_server_ws_url

        url = build_server_ws_url(self.server_url)

        headers = {}
        if self.token:
            headers["Authorization"] = f"Bearer {self.token}"
        if self.client_id:
            headers["X-Client-Id"] = self.client_id

        try:
            self.ws = await ws_connect(
                url,
                additional_headers=headers,
                ping_interval=None,
                ping_timeout=None,
            )
        except (OSError, WebSocketException) as e:
            raise Exception(f"Failed to connect to server at {self.server_url}: {e}")

        welcome = await self._receive()
        if welcome.get("type") == "error":
            msg = welcome.get("message", "Authentication failed")
            await self.ws.close()
            self.ws = None
            raise Exception(f"Server rejected connection: {msg}")

        if welcome.get("type") != "welcome":
            raise Exception(f"Expected welcome, got: {welcome.get('type')}")

        self.client_id = welcome.get("client_id", self.client_id)
        self.connected = True
        return self.client_id

    async def register(self, gateway_info: dict | None = None, resources: dict | None = None):
        """Send registration message with capabilities."""
        if not self.ws or not self.connected:
            raise Exception("Not connected to server")

        msg = {
            "type": "register",
            "client_info": {
                "hostname": self.hostname,
                "version": "0.1.0",
                "platform": platform.system(),
            },
            "gateway": gateway_info or {"connected": False},
            "resources": resources or {},
        }
        await self.ws.send(json.dumps(msg))

        response = await self._receive()
        if response.get("type") != "registered":
            raise Exception(f"Registration failed: {response}")

        logger.info("Registered with server as %s", self.client_id)

    async def start_heartbeat(self):
        """Start the heartbeat loop."""
        self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())

    async def _heartbeat_loop(self):
        while self.connected:
            await asyncio.sleep(HEARTBEAT_INTERVAL)
            if not self.connected or not self.ws:
                break
            try:
                await self.ws.send(json.dumps({"type": "heartbeat"}))
            except Exception:
                self.connected = False
                break

    async def send_chat_event(self, request_id: str, payload: dict):
        """Send a streaming chat event back to the server."""
        if not self.ws or not self.connected:
            return
        msg = {
            "type": "chat_event",
            "request_id": request_id,
            "payload": payload,
        }
        await self.ws.send(json.dumps(msg))

    async def send_command_result(self, request_id: str, ok: bool, result=None, error=None):
        """Send a command result back to the server."""
        if not self.ws or not self.connected:
            return
        msg = {
            "type": "command_result",
            "id": request_id,
            "ok": ok,
        }
        if result is not None:
            msg["result"] = result
        if error is not None:
            msg["error"] = error
        await self.ws.send(json.dumps(msg))

    async def send_gateway_status(self, connected: bool):
        """Notify server about gateway connection status change."""
        if not self.ws or not self.connected:
            return
        await self.ws.send(json.dumps({
            "type": "gateway_status",
            "connected": connected,
        }))

    async def listen(self, on_command: Callable[[dict], Awaitable]):
        """Listen for incoming messages from the server."""
        self._on_command = on_command
        try:
            async for raw in self.ws:
                data = json.loads(raw)
                msg_type = data.get("type")

                if msg_type == "heartbeat_ack":
                    continue
                elif msg_type == "command":
                    asyncio.create_task(on_command(data))
                elif msg_type == "error":
                    logger.warning("Server error: %s", data.get("message"))
        except ConnectionClosed:
            pass
        finally:
            self.connected = False

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

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

    async def reconnect_loop(self) -> str:
        """Reconnect with exponential backoff. Returns client_id on success."""
        backoff = 1
        while True:
            logger.info("Attempting server reconnect in %ds...", backoff)
            await asyncio.sleep(backoff)
            try:
                client_id = await self.connect()
                logger.info("Server reconnected as %s", client_id)
                return client_id
            except Exception as e:
                logger.warning("Server reconnect failed: %s", e)
                backoff = min(backoff * 2, MAX_BACKOFF)
