"""WebSocket keyword library for Robot Framework tests."""

import asyncio
from typing import Optional

import websockets
from robot.api import logger


class WebSocketKeywords:
    """Keywords for WebSocket testing."""

    def __init__(self):
        self._connection: Optional[websockets.WebSocketClientProtocol] = None

    def connect_to_websocket(self, ws_url: str) -> str:
        """Establish a WebSocket connection and return connection info.

        Args:
            ws_url: Full WebSocket URL including token query param.

        Returns:
            String representation of the connected WebSocket.
        """
        asyncio.run(self._async_connect(ws_url))
        return f"Connected to {ws_url}"

    async def _async_connect(self, ws_url: str) -> None:
        self._connection = await websockets.connect(ws_url)
        logger.info(f"WebSocket connected: {self._connection}")

    def send_websocket_message(self, message: str) -> None:
        """Send a text message over the active WebSocket connection.

        Args:
            message: Text message to send.
        """
        if self._connection is None:
            raise RuntimeError("No active WebSocket connection. Use Connect To WebSocket first.")
        asyncio.run(self._connection.send(message))
        logger.info(f"Sent message: {message}")

    def receive_websocket_message(self, timeout: float = 5.0) -> str:
        """Receive a text message from the active WebSocket connection.

        Args:
            timeout: Maximum seconds to wait for a message.

        Returns:
            The received text message.
        """
        if self._connection is None:
            raise RuntimeError("No active WebSocket connection. Use Connect To WebSocket first.")
        try:
            message = asyncio.run(asyncio.wait_for(self._connection.recv(), timeout=timeout))
            logger.info(f"Received message: {message}")
            return message
        except asyncio.TimeoutError:
            raise TimeoutError(f"No message received within {timeout} seconds")

    def close_websocket(self) -> None:
        """Close the active WebSocket connection."""
        if self._connection is not None:
            asyncio.run(self._connection.close())
            self._connection = None
            logger.info("WebSocket closed")