"""
WebSocket Client for Browse-Now CLI

Connects to FastAPI server's WebSocket endpoint to control browser via extension.
"""

import asyncio
import json
import sys
import uuid
from typing import Any, Optional

from pydantic import BaseModel


class WSMessage(BaseModel):
    """WebSocket message format."""

    id: str
    type: str
    browser_id: Optional[str] = None
    payload: Any = None


class WSResponse(BaseModel):
    """WebSocket response format."""

    id: str
    success: bool
    data: Any = None
    error: Optional[str] = None


class BrowserInfo(BaseModel):
    """Browser information."""

    browser_id: str
    browser_name: str
    user_agent: str
    extension_version: str
    connected_at: str


class BrowsersResponse(BaseModel):
    """List browsers response."""

    browsers: list[BrowserInfo]
    count: int


class BrowseNowWSClient:
    """
    WebSocket client for browse-now CLI.

    Connects to FastAPI server to send commands to browser extension.
    """

    def __init__(
        self,
        ws_url: str = "ws://127.0.0.1:14242/browser-bridge/ws",
        browser_id: Optional[str] = None,
        timeout: float = 30.0,
    ):
        self.ws_url = ws_url
        self.browser_id = browser_id
        self.timeout = timeout
        self._ws: Any = None  # websockets.WebSocketClientProtocol
        self._connected = False

    async def connect(self) -> None:
        """Connect to WebSocket server."""
        try:
            import websockets

            # Set close_timeout to handle servers that don't send close frames properly
            # max_size increased to 10MB for large screenshot responses (default is 1MB)
            self._ws = await websockets.connect(
                self.ws_url,
                close_timeout=2,  # Don't wait long for close frame
                ping_timeout=10,
                ping_interval=None,  # Disable pings for short-lived connections
                max_size=10 * 1024 * 1024,  # 10MB - screenshots can be 2-5MB
            )
            self._connected = True
        except ImportError:
            raise ImportError("websockets package required. Install with: pip install websockets")

    async def disconnect(self) -> None:
        """Disconnect from WebSocket server."""
        if self._ws:
            try:
                await self._ws.close()
            except Exception:
                pass  # Ignore close errors - connection may already be closed
            finally:
                self._ws = None
                self._connected = False

    async def __aenter__(self) -> "BrowseNowWSClient":
        await self.connect()
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.disconnect()

    async def _send_and_receive(self, message_type: str, payload: Any = None) -> WSResponse:
        """Send message and wait for response."""
        if not self._ws:
            raise RuntimeError("Not connected. Call connect() first.")

        message_id = str(uuid.uuid4())
        message = WSMessage(
            id=message_id, type=message_type, browser_id=self.browser_id, payload=payload
        )

        try:
            await self._ws.send(message.model_dump_json())

            # Wait for response with timeout
            response_data = await asyncio.wait_for(self._ws.recv(), timeout=self.timeout)
            response_dict = json.loads(response_data)
            return WSResponse(**response_dict)
        except asyncio.TimeoutError:
            return WSResponse(
                id=message_id, success=False, error=f"Request timed out after {self.timeout}s"
            )
        except Exception as e:
            # Handle WebSocket close errors gracefully
            error_msg = str(e)
            if "no close frame" in error_msg.lower():
                return WSResponse(
                    id=message_id, success=False, error="Connection closed unexpectedly"
                )
            raise

    # =========================================================================
    # Browser Management
    # =========================================================================

    async def list_browsers(self) -> BrowsersResponse:
        """List all connected browsers."""
        response = await self._send_and_receive("list_browsers")
        if response.success and response.data:
            try:
                return BrowsersResponse(**response.data)
            except Exception:
                # Data parsing failed - return empty
                return BrowsersResponse(browsers=[], count=0)
        # Return empty response with error info preserved for debugging
        return BrowsersResponse(browsers=[], count=0)

    async def ping(self) -> bool:
        """Ping the server."""
        response = await self._send_and_receive("ping")
        return response.success and response.data == "pong"

    # =========================================================================
    # Browser Commands
    # =========================================================================

    async def send_command(
        self, action: str, params: Optional[dict[str, Any]] = None
    ) -> WSResponse:
        """Send a command to the browser."""
        payload = {"action": action, "params": params or {}}
        return await self._send_and_receive("command", payload)

    async def snapshot(
        self,
        interactive: bool = True,
        max_depth: int = 10,
        include_coordinates: bool = False,
        selector: str | None = None,
        with_css: bool = False,
        wait_ms: int = 0,
    ) -> WSResponse:
        """Capture page snapshot with element refs.

        Args:
            interactive: Only include interactive elements
            max_depth: Maximum depth to traverse
            include_coordinates: Include element coordinates
            selector: Scope snapshot to CSS selector (e.g., '#main-content')
            with_css: Include CSS selectors for each element
            wait_ms: Wait milliseconds before snapshot (for late-loading content)
        """
        params: dict[str, Any] = {
            "interactive": interactive,
            "maxDepth": max_depth,
            "includeCoordinates": include_coordinates,
            "withCss": with_css,
        }
        if selector:
            params["selector"] = selector
        if wait_ms > 0:
            params["waitMs"] = wait_ms
        return await self.send_command("snapshot", params)

    async def click(self, ref: str) -> WSResponse:
        """Click an element by ref (e.g., @e1 or e1)."""
        return await self.send_command("click", {"ref": ref})

    async def fill(
        self, ref: str, text: str, clear: bool = True, submit: bool = False
    ) -> WSResponse:
        """Fill an input element by ref."""
        return await self.send_command(
            "fill", {"ref": ref, "text": text, "clear": clear, "submit": submit}
        )

    async def navigate(
        self, url: str, wait_until: str = "load", use_agent_tab: bool = False, timeout: int = 10
    ) -> WSResponse:
        """Navigate to a URL.

        Args:
            url: The URL to navigate to
            wait_until: When to consider navigation complete
            use_agent_tab: If True, opens in isolated "Browse-Now" tab group
            timeout: Page load timeout in seconds (default 10)
        """
        return await self.send_command(
            "navigate",
            {
                "url": url,
                "waitUntil": wait_until,
                "useAgentTab": use_agent_tab,
                "timeout": timeout * 1000,  # Convert to milliseconds for extension
            },
        )

    async def tabs_list(self) -> WSResponse:
        """List open browser tabs."""
        return await self.send_command("tabs-list")

    async def tabs_switch(self, tab_id: int) -> WSResponse:
        """Switch to a specific tab."""
        return await self.send_command("tabs-switch", {"tabId": tab_id})

    async def status(self) -> WSResponse:
        """Get browser bridge status."""
        return await self.send_command("status")


class BrowseNowWSClientSync:
    """Synchronous wrapper for BrowseNowWSClient."""

    def __init__(
        self,
        ws_url: str = "ws://127.0.0.1:14242/browser-bridge/ws",
        browser_id: Optional[str] = None,
        timeout: float = 30.0,
    ):
        self._ws_url = ws_url
        self._browser_id = browser_id
        self._timeout = timeout

    def _run(self, coro: Any) -> Any:
        """Run an async coroutine synchronously."""
        return asyncio.run(coro)

    def _with_client(self, method_name: str, *args: Any, **kwargs: Any) -> Any:
        """Execute method with a fresh client connection.

        Retries up to 3 times with 2s delay when browser is not yet connected
        (e.g., after app restart while extension reconnects via keepalive alarm).
        """
        max_retries = 3
        retry_delay = 2.0  # seconds

        async def _execute():
            async with BrowseNowWSClient(
                ws_url=self._ws_url, browser_id=self._browser_id, timeout=self._timeout
            ) as client:
                method = getattr(client, method_name)
                return await method(*args, **kwargs)

        for attempt in range(max_retries + 1):
            result = self._run(_execute())

            # Check if we got "No browser connected" and should retry
            if (
                attempt < max_retries
                and isinstance(result, WSResponse)
                and not result.success
                and result.error
                and "No browser connected" in result.error
            ):
                print(
                    f"⏳ No browser connected, waiting for extension reconnect... "
                    f"(retry {attempt + 1}/{max_retries})",
                    file=sys.stderr,
                )
                self._run(asyncio.sleep(retry_delay))
                continue

            return result

        return result  # type: ignore[possibly-undefined]

    # Browser Management
    def list_browsers(self) -> BrowsersResponse:
        return self._with_client("list_browsers")

    def ping(self) -> bool:
        return self._with_client("ping")

    # Browser Commands
    def send_command(self, action: str, params: Optional[dict[str, Any]] = None) -> WSResponse:
        return self._with_client("send_command", action, params)

    def snapshot(
        self,
        interactive: bool = True,
        max_depth: int = 10,
        include_coordinates: bool = False,
        selector: str | None = None,
        with_css: bool = False,
        wait_ms: int = 0,
    ) -> WSResponse:
        return self._with_client(
            "snapshot", interactive, max_depth, include_coordinates, selector, with_css, wait_ms
        )

    def click(self, ref: str) -> WSResponse:
        return self._with_client("click", ref)

    def fill(self, ref: str, text: str, clear: bool = True, submit: bool = False) -> WSResponse:
        return self._with_client("fill", ref, text, clear, submit)

    def navigate(
        self, url: str, wait_until: str = "load", use_agent_tab: bool = False, timeout: int = 10
    ) -> WSResponse:
        return self._with_client("navigate", url, wait_until, use_agent_tab, timeout)

    def tabs_list(self) -> WSResponse:
        return self._with_client("tabs_list")

    def tabs_switch(self, tab_id: int) -> WSResponse:
        return self._with_client("tabs_switch", tab_id)

    def status(self) -> WSResponse:
        return self._with_client("status")
