"""
WebSocket tunnel client for the local bridge daemon.

Opens an outbound WSS connection to the production relay, authenticates,
registers available servers, and handles bidirectional message forwarding.
Reconnects automatically with exponential backoff on disconnection.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import signal
from typing import Optional

import websockets
from websockets.asyncio.client import ClientConnection

from .config import BridgeConfig
from .protocol import (
    MSG_AUTH_ERROR,
    MSG_AUTH_OK,
    MSG_ERROR,
    MSG_HEARTBEAT_ACK,
    MSG_MCP_REQUEST,
    HEARTBEAT_INTERVAL_SEC,
    build_auth_message,
    build_heartbeat_message,
    build_mcp_response_message,
    build_register_servers_message,
)
from .server_manager import ServerManager

logger = logging.getLogger(__name__)

# Reconnection backoff
RECONNECT_BASE_SEC = float(os.environ.get("NEXUS_BRIDGE_RECONNECT_BASE_SEC", "1.0"))
RECONNECT_MAX_SEC = float(os.environ.get("NEXUS_BRIDGE_RECONNECT_MAX_SEC", "60.0"))
if RECONNECT_BASE_SEC <= 0:
    raise ValueError("NEXUS_BRIDGE_RECONNECT_BASE_SEC must be > 0.")
if RECONNECT_MAX_SEC < RECONNECT_BASE_SEC:
    raise ValueError(
        "NEXUS_BRIDGE_RECONNECT_MAX_SEC must be >= NEXUS_BRIDGE_RECONNECT_BASE_SEC."
    )


class BridgeTunnel:
    """
    Manages the WebSocket tunnel to the production relay.

    Lifecycle:
    1. Connect to relay WSS endpoint
    2. Authenticate with JWT
    3. Register available servers
    4. Forward MCP requests to local servers, send responses back
    5. Heartbeat keepalive
    6. Reconnect on disconnect
    """

    def __init__(
        self,
        relay_url: str,
        token: str,
        server_manager: ServerManager,
    ):
        self._relay_url = relay_url
        self._token = token
        self._server_manager = server_manager
        self._ws: Optional[ClientConnection] = None
        self._running = False
        self._heartbeat_task: Optional[asyncio.Task] = None
        self._heartbeat_interval_sec = HEARTBEAT_INTERVAL_SEC

    async def run(self) -> None:
        """Run the tunnel with automatic reconnection."""
        self._running = True
        attempt = 0

        while self._running:
            try:
                await self._connect_and_serve()
                # Clean disconnect; reset backoff
                attempt = 0
            except asyncio.CancelledError:
                break
            except Exception as exc:
                if not self._running:
                    break
                attempt += 1
                backoff = min(RECONNECT_BASE_SEC * (2 ** (attempt - 1)), RECONNECT_MAX_SEC)
                logger.warning(
                    "[Tunnel] Connection lost (%s). Reconnecting in %.1fs (attempt %d)...",
                    exc, backoff, attempt,
                )
                await asyncio.sleep(backoff)

        await self._server_manager.stop_all()
        logger.info("[Tunnel] Shutdown complete.")

    async def stop(self) -> None:
        """Initiate graceful shutdown."""
        self._running = False
        if self._heartbeat_task:
            self._heartbeat_task.cancel()
        if self._ws:
            await self._ws.close()

    async def _connect_and_serve(self) -> None:
        """Single connection lifecycle."""
        logger.info("[Tunnel] Connecting to %s ...", self._relay_url)

        async with websockets.connect(
            self._relay_url,
            additional_headers={"Authorization": f"Bearer {self._token}"},
            ping_interval=20,
            ping_timeout=10,
            close_timeout=5,
        ) as ws:
            self._ws = ws

            # Authenticate
            await ws.send(json.dumps(build_auth_message(self._token)))
            auth_resp = json.loads(await ws.recv())

            if auth_resp.get("type") == MSG_AUTH_ERROR:
                raise RuntimeError(f"Authentication failed: {auth_resp.get('message')}")
            if auth_resp.get("type") != MSG_AUTH_OK:
                raise RuntimeError(f"Unexpected auth response: {auth_resp}")
            heartbeat_interval = auth_resp.get("heartbeat_interval_sec")
            if isinstance(heartbeat_interval, int) and heartbeat_interval > 0:
                self._heartbeat_interval_sec = heartbeat_interval

            session_id = auth_resp.get("session_id")
            logger.info("[Tunnel] Authenticated. Session: %s", session_id)

            # Register servers
            manifest = self._server_manager.get_manifest()
            await ws.send(json.dumps(build_register_servers_message(manifest)))
            logger.info(
                "[Tunnel] Registered %d server(s): %s",
                len(manifest), [s["name"] for s in manifest],
            )

            # Start heartbeat
            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop(ws))

            # Message loop
            try:
                async for raw in ws:
                    msg = json.loads(raw)
                    await self._handle_message(ws, msg)
            finally:
                if self._heartbeat_task:
                    self._heartbeat_task.cancel()
                    try:
                        await self._heartbeat_task
                    except asyncio.CancelledError:
                        pass

    async def _handle_message(self, ws: ClientConnection, msg: dict) -> None:
        """Handle a single message from the relay."""
        msg_type = msg.get("type")

        if msg_type == MSG_MCP_REQUEST:
            request_id = msg.get("request_id")
            server_name = msg.get("server_name")
            payload = msg.get("payload")

            if not request_id or not server_name or not isinstance(payload, dict):
                logger.warning("[Tunnel] Invalid mcp_request: %s", msg)
                return

            # Forward to local server
            response = await self._server_manager.handle_request(server_name, payload)

            # Send response back to relay
            resp_msg = build_mcp_response_message(request_id, server_name, response)
            await ws.send(json.dumps(resp_msg))

        elif msg_type == MSG_HEARTBEAT_ACK:
            pass  # Expected response to our heartbeat

        elif msg_type == MSG_ERROR:
            logger.warning(
                "[Tunnel] Server error: %s - %s",
                msg.get("code"), msg.get("message"),
            )

        else:
            logger.debug("[Tunnel] Unknown message type: %s", msg_type)

    async def _heartbeat_loop(self, ws: ClientConnection) -> None:
        """Send periodic heartbeats to keep the connection alive."""
        try:
            while self._running:
                await asyncio.sleep(self._heartbeat_interval_sec)
                if not self._running:
                    break
                await ws.send(json.dumps(build_heartbeat_message()))
        except asyncio.CancelledError:
            pass
        except Exception as exc:
            logger.debug("[Tunnel] Heartbeat stopped: %s", exc)
