"""
WebSocket Adapter — Expose a WebSocket endpoint as Brynq tools.

Accepts a WebSocket URL (``ws://…`` or ``wss://…``) and provides tools
to connect, send messages, and listen for incoming events.

Uses ``asyncio`` and ``websockets`` (lazy imported) for connection
management.  All stdlib except ``websockets`` which is imported lazily.
"""

from __future__ import annotations

from typing import Any, Dict, List

from runtime.adapters.base import BaseAdapter

_VERSION = "0.1.0"


class WebSocketAdapter(BaseAdapter):
    """Adapter that exposes a WebSocket endpoint as callable tools.

    Parameters
    ----------
    source:
        A WebSocket URL (e.g. ``ws://localhost:8080/ws`` or
        ``wss://api.example.com/stream``).
    """

    def __init__(self, source: str) -> None:
        self._source = source
        self._connection: Any = None
        self._tools: List[Dict[str, Any]] | None = None

    # ── BaseAdapter interface ────────────────────────────────────────────

    def introspect(self) -> Dict[str, Any]:
        """Return the fixed set of WebSocket tools."""
        self._ensure_tools()
        return {
            "tools": list(self._tools),  # type: ignore[arg-type]
            "version": _VERSION,
            "source": self._source,
        }

    def execute(self, tool_name: str, args: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a WebSocket tool by name.

        Supported tools:
          - ``connect``      — Open a connection to the WebSocket URL.
          - ``send_message`` — Send a message over the open connection.
          - ``listen``       — Receive the next N messages from the connection.

        Parameters
        ----------
        tool_name:
            Name of the tool to invoke.
        args:
            Keyword arguments for the operation.
        """
        self._ensure_tools()

        tool = self._find_tool(tool_name)
        if tool is None:
            return {"status": "error", "result": f"Unknown tool: {tool_name}"}

        try:
            result = self._run_tool(tool, args)
            return {"status": "ok", "result": result}
        except Exception as exc:  # noqa: BLE001
            return {"status": "error", "result": str(exc)}

    # ── Tool setup ───────────────────────────────────────────────────────

    def _ensure_tools(self) -> None:
        """Build the tool list lazily (once)."""
        if self._tools is not None:
            return
        self._tools = [
            {
                "name": "connect",
                "description": f"Open a WebSocket connection to {self._source}",
                "parameters": [],
            },
            {
                "name": "send_message",
                "description": "Send a message over the open WebSocket connection",
                "parameters": [
                    {"name": "message", "description": "The message to send", "required": True},
                ],
            },
            {
                "name": "listen",
                "description": "Receive the next N messages from the WebSocket",
                "parameters": [
                    {"name": "count", "description": "Number of messages to receive (default: 1)", "required": False},
                    {"name": "timeout", "description": "Timeout in seconds (default: 30)", "required": False},
                ],
            },
        ]

    def _find_tool(self, tool_name: str) -> Dict[str, Any] | None:
        """Look up a tool by name."""
        assert self._tools is not None  # noqa: S101
        for tool in self._tools:
            if tool["name"] == tool_name:
                return tool
        return None

    # ── Execution ────────────────────────────────────────────────────────

    def _run_tool(self, tool: Dict[str, Any], args: Dict[str, Any]) -> Any:
        """Dispatch to the appropriate WebSocket operation."""
        import asyncio  # noqa: E402 — deferred to avoid import cost

        name = tool["name"]

        if name == "connect":
            return asyncio.get_event_loop().run_until_complete(self._connect())
        if name == "send_message":
            message = args.get("message")
            if not message:
                raise ValueError("Missing required parameter: message")
            return asyncio.get_event_loop().run_until_complete(self._send(message))
        if name == "listen":
            count = int(args.get("count", 1))
            timeout = float(args.get("timeout", 30))
            return asyncio.get_event_loop().run_until_complete(
                self._listen(count, timeout)
            )

        raise ValueError(f"Unknown operation: {name}")

    # ── Async WebSocket operations ───────────────────────────────────────

    async def _connect(self) -> Dict[str, Any]:
        """Open a WebSocket connection."""
        websockets = _import_websockets()
        self._connection = await websockets.connect(self._source)
        return {"connected": True, "url": self._source}

    async def _send(self, message: str) -> Dict[str, Any]:
        """Send a message over the open connection."""
        if self._connection is None:
            raise RuntimeError("Not connected. Call 'connect' first.")
        await self._connection.send(message)
        return {"sent": True, "message": message}

    async def _listen(self, count: int, timeout: float) -> List[str]:
        """Receive the next *count* messages, with a per-message timeout."""
        import asyncio  # noqa: E402

        if self._connection is None:
            raise RuntimeError("Not connected. Call 'connect' first.")

        messages: List[str] = []
        for _ in range(count):
            try:
                msg = await asyncio.wait_for(
                    self._connection.recv(), timeout=timeout
                )
                messages.append(msg)
            except asyncio.TimeoutError:
                break
        return messages


# ── helpers ──────────────────────────────────────────────────────────────


def _import_websockets():  # type: ignore[no-untyped-def]
    """Lazily import websockets."""
    try:
        import websockets  # noqa: E402
        return websockets
    except ImportError:
        raise ImportError(
            "websockets is required for WebSocket connections. "
            "Install it with: pip install websockets"
        ) from None
