"""
🔌 WebSocket client for Moltbot/Clawdbot communication.

Implements the Clawdbot Gateway protocol v3.
Forwards your couch wisdom to the bot.
"""

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

import websockets
from websockets.client import WebSocketClientProtocol

logger = logging.getLogger(__name__)

# Default Moltbot WebSocket URL
DEFAULT_WS_URL = "ws://127.0.0.1:18789"

# BlueLobster version
VERSION = "0.1.0"

# Protocol version
PROTOCOL_VERSION = 3


class MoltbotClient:
    """
    WebSocket client for communicating with Moltbot/Clawdbot.
    
    Implements the Clawdbot Gateway protocol v3 with proper handshake.
    """
    
    def __init__(
        self,
        ws_url: str = DEFAULT_WS_URL,
        token: Optional[str] = None,
        on_response: Optional[Callable[[str], Awaitable[None]]] = None,
    ):
        """
        Initialize the Moltbot client.
        
        Args:
            ws_url: WebSocket URL for Moltbot gateway.
            token: Optional authentication token.
            on_response: Optional async callback for bot responses.
        """
        self.ws_url = ws_url
        self.token = token
        self.on_response = on_response
        self._ws: Optional[WebSocketClientProtocol] = None
        self._connected = False
        self._handshake_complete = False
        self._listen_task: Optional[asyncio.Task] = None
        self._request_id = 0
        self._instance_id = str(uuid.uuid4())
        self._session_key: Optional[str] = None
    
    @property
    def connected(self) -> bool:
        """Check if currently connected and handshake complete."""
        return self._connected and self._handshake_complete and self._ws is not None
    
    def _next_request_id(self) -> str:
        """Generate next request ID."""
        self._request_id += 1
        return f"bluelobster-{self._request_id}"
    
    async def connect(self) -> bool:
        """
        Connect to the Moltbot WebSocket and perform handshake.
        
        Returns:
            True if connection and handshake successful, False otherwise.
        """
        try:
            logger.info(f"🔌 Connecting to Moltbot at {self.ws_url}...")
            self._ws = await websockets.connect(self.ws_url)
            self._connected = True
            
            # Perform the mandatory connect handshake
            handshake_ok = await self._perform_handshake()
            if not handshake_ok:
                logger.error("❌ Handshake failed")
                await self._ws.close()
                self._connected = False
                return False
            
            self._handshake_complete = True
            logger.info("✅ Connected to Moltbot!")
            
            # Start listening for responses if callback provided
            if self.on_response:
                self._listen_task = asyncio.create_task(self._listen_loop())
            
            return True
        except Exception as e:
            logger.error(f"❌ Failed to connect to Moltbot: {e}")
            self._connected = False
            return False
    
    async def _perform_handshake(self) -> bool:
        """
        Perform the Clawdbot gateway handshake.
        
        1. Receive connect.challenge from gateway
        2. Send connect request with auth and client info
        3. Receive hello-ok response
        """
        try:
            # Wait for challenge
            challenge_raw = await asyncio.wait_for(self._ws.recv(), timeout=10.0)
            challenge = json.loads(challenge_raw)
            
            if challenge.get("event") != "connect.challenge":
                logger.error(f"❌ Expected connect.challenge, got: {challenge}")
                return False
            
            logger.debug("📥 Received challenge")
            
            # Build connect request
            connect_request = {
                "type": "req",
                "id": self._next_request_id(),
                "method": "connect",
                "params": {
                    "minProtocol": PROTOCOL_VERSION,
                    "maxProtocol": PROTOCOL_VERSION,
                    "client": {
                        "id": "cli",
                        "displayName": "BlueLobster",
                        "version": VERSION,
                        "platform": platform.system().lower(),
                        "mode": "webchat",  # Connect as webchat client
                    },
                    "role": "operator",
                    "scopes": ["operator.read", "operator.write"],
                    "caps": [],
                    "commands": [],
                    "permissions": {},
                    "userAgent": f"BlueLobster/{VERSION} Python/{platform.python_version()}",
                }
            }
            
            # Add auth token if provided
            if self.token:
                connect_request["params"]["auth"] = {"token": self.token}
            
            await self._ws.send(json.dumps(connect_request))
            logger.debug("📤 Sent connect request")
            
            # Wait for response
            response_raw = await asyncio.wait_for(self._ws.recv(), timeout=10.0)
            response = json.loads(response_raw)
            
            if response.get("type") == "res" and response.get("ok"):
                # Extract session info
                payload = response.get("payload", {})
                snapshot = payload.get("snapshot", {})
                session_defaults = snapshot.get("sessionDefaults", {})
                self._session_key = session_defaults.get("mainSessionKey", "agent:main:main")
                logger.debug(f"✅ Handshake OK, session: {self._session_key}")
                return True
            else:
                error = response.get("error", {})
                logger.error(f"❌ Handshake rejected: {error}")
                return False
        except asyncio.TimeoutError:
            logger.error("❌ Handshake timeout")
            return False
        except Exception as e:
            logger.error(f"❌ Handshake error: {e}")
            return False
    
    async def disconnect(self) -> None:
        """Disconnect from Moltbot."""
        self._connected = False
        self._handshake_complete = False
        
        if self._listen_task:
            self._listen_task.cancel()
            try:
                await self._listen_task
            except asyncio.CancelledError:
                pass
            self._listen_task = None
        
        if self._ws:
            await self._ws.close()
            self._ws = None
            logger.info("🔌 Disconnected from Moltbot")
    
    async def send_message(self, text: str, sender: str = "BlueLobster") -> bool:
        """
        Send a message to Moltbot/Clawdbot.
        
        Uses the Clawdbot chat.send method.
        
        Args:
            text: The message text to send.
            sender: Sender identifier (default: "BlueLobster").
        
        Returns:
            True if message sent successfully, False otherwise.
        """
        if not self.connected:
            logger.warning("⚠️ Not connected to Moltbot, attempting to connect...")
            if not await self.connect():
                return False
        
        # chat.send requires sessionKey and idempotencyKey
        message = {
            "type": "req",
            "id": self._next_request_id(),
            "method": "chat.send",
            "params": {
                "message": text,
                "sessionKey": self._session_key or "agent:main:main",
                "idempotencyKey": str(uuid.uuid4()),
            }
        }
        
        try:
            await self._ws.send(json.dumps(message))
            logger.info(f"📤 Sent to Moltbot: {text[:50]}{'...' if len(text) > 50 else ''}")
            return True
        except Exception as e:
            logger.error(f"❌ Failed to send message: {e}")
            self._connected = False
            return False
    
    async def _listen_loop(self) -> None:
        """Listen for responses from Moltbot."""
        try:
            while self._connected and self._ws:
                try:
                    message = await self._ws.recv()
                    
                    # Parse and handle different message types
                    try:
                        data = json.loads(message)
                        msg_type = data.get("type")
                        
                        if msg_type == "event":
                            event_name = data.get("event")
                            if event_name == "agent":
                                # Agent streaming output
                                payload = data.get("payload", {})
                                content = payload.get("content") or payload.get("text") or payload.get("delta")
                                if content:
                                    logger.debug(f"🤖 Agent: {content[:100]}")
                        elif msg_type == "res":
                            # Response to a request
                            if not data.get("ok"):
                                error = data.get("error", {})
                                logger.warning(f"⚠️ Request failed: {error}")
                    except json.JSONDecodeError:
                        pass
                    
                    if self.on_response:
                        await self.on_response(message)
                except websockets.ConnectionClosed:
                    logger.warning("🔌 Moltbot connection closed")
                    self._connected = False
                    break
        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.error(f"Error in listen loop: {e}")
            self._connected = False


async def send_to_moltbot(
    text: str,
    ws_url: str = DEFAULT_WS_URL,
    token: Optional[str] = None,
) -> bool:
    """
    One-shot function to send a message to Moltbot.
    
    Args:
        text: Message text to send.
        ws_url: WebSocket URL for Moltbot.
        token: Optional authentication token.
    
    Returns:
        True if sent successfully, False otherwise.
    """
    client = MoltbotClient(ws_url=ws_url, token=token)
    try:
        if await client.connect():
            return await client.send_message(text)
        return False
    finally:
        await client.disconnect()
