"""WebSocket client for Runner to Backend communication."""

import asyncio
import json
import logging
import random
from typing import Optional

import websockets
import httpx
from websockets.exceptions import ConnectionClosed, ConnectionClosedError

from .config import RunnerConfig
from .message_handler import MessageHandler
from .output_buffer import OutputBuffer, BUFFERED_METHODS
from .rate_limiter import SlidingWindowRateLimiter
from .structured_log import get_logger as _get_slog

logger = logging.getLogger("skop-runner")
_log = _get_slog("ws_client")


# Rate limiting for incoming messages (prevents DoS from malicious backend)
MAX_INCOMING_PER_SECOND = 200
INCOMING_WINDOW_SECONDS = 1.0

# WebSocket connection settings
WS_PING_INTERVAL = 15  # Faster dead connection detection (was 30)
WS_PING_TIMEOUT = 10


class RunnerClient:
    """WebSocket client for Runner ↔ Backend communication."""

    def __init__(self, config: RunnerConfig, message_handler: MessageHandler, output_buffer: OutputBuffer):
        self.config = config
        self.message_handler = message_handler
        self._output_buffer = output_buffer
        self.ws: Optional[websockets.WebSocketClientProtocol] = None
        self.running = False
        self._reconnect_delay = 1  # Initial reconnect delay in seconds
        self._max_reconnect_delay = 15  # Max reconnect delay
        self.on_connect: Optional[callable] = None  # Callback for successful connection
        self._should_reconnect = True
        # Rate limiting for incoming messages (prevents DoS from relay)
        self._incoming_limiter = SlidingWindowRateLimiter(
            MAX_INCOMING_PER_SECOND, INCOMING_WINDOW_SECONDS,
            label="incoming relay messages",
        )
        # Shared HTTP client for heartbeats (created lazily, closed on disconnect)
        self._http_client: Optional[httpx.AsyncClient] = None
        # Lock for serializing WebSocket sends from concurrent handlers
        self._send_lock = asyncio.Lock()
        # Optional callable returning metadata dict for heartbeat payload
        self._metadata_supplier: Optional[callable] = None

    @property
    def output_buffer(self) -> OutputBuffer:
        """Access the output buffer."""
        return self._output_buffer

    async def connect(self) -> bool:
        """
        Connect to WebSocket server.

        Returns:
            True if connected successfully, False otherwise
        """
        logger.info(f"Connecting to WebSocket server...")

        try:
            # Token is sent in URL query string as required by backend
            ws_url = f"{self.config.ws_url}?token={self.config.device_token}"
            self.ws = await websockets.connect(
                ws_url,
                ping_interval=WS_PING_INTERVAL,
                ping_timeout=WS_PING_TIMEOUT,
                max_size=50 * 1024 * 1024,  # 50MB — match local WS server limit
            )

            self.running = True
            self._output_buffer.set_connected(True)
            self._reconnect_delay = 1  # Reset on successful connection
            self._incoming_limiter.reset()

            _log.info("ws.connect", url=self.config.ws_url)
            return True

        except Exception as e:
            logger.error(f"Failed to connect: {e}")
            return False

    async def _send_notification(self, method: str, params: dict):
        """
        Send a JSON-RPC notification (no id, no response expected).

        No outgoing rate limiting — natural backpressure from awaited WS
        sends throttles the caller. Buffers to OutputBuffer when disconnected.
        """
        if not self.ws or not self.running:
            if method in BUFFERED_METHODS:
                self._output_buffer.add(method, params)
            return

        notification = {
            "jsonrpc": "2.0",
            "method": method,
            "params": params,
        }
        try:
            async with self._send_lock:
                await self.ws.send(json.dumps(notification))
            logger.debug(f"Sent notification: {method}")
        except Exception as e:
            logger.warning(f"Failed to send notification {method}: {e}")
            if method in BUFFERED_METHODS:
                self._output_buffer.add(method, params)

    async def disconnect(self):
        """Disconnect from server."""
        self._output_buffer.set_connected(False)
        self.running = False
        self._should_reconnect = True

        # Cancel any running tasks
        if hasattr(self, "_tasks"):
            for task in self._tasks:
                if not task.done():
                    task.cancel()
            # Wait for tasks to finish cancelling
            for task in self._tasks:
                try:
                    await task
                except asyncio.CancelledError:
                    pass

        if self.ws:
            try:
                await self.ws.close()
            except Exception as e:
                logger.warning(f"Error closing WebSocket: {e}")

        # Close shared HTTP client
        if self._http_client:
            try:
                await self._http_client.aclose()
            except Exception as e:
                logger.warning(f"Error closing HTTP client: {e}")
            self._http_client = None

        _log.info("ws.disconnect")

    async def _handle_and_respond(self, message: dict):
        """Handle a message and send the response. Runs as a concurrent task."""
        try:
            response = await self.message_handler.handle_message(message)
            if self.ws and self.running:
                async with self._send_lock:
                    await self.ws.send(json.dumps(response))
                logger.debug(f"Sent response for id: {response.get('id')}")
        except Exception as e:
            logger.error(f"Error handling message: {e}", exc_info=True)

    async def receive_loop(self):
        """Receive and handle messages from server."""
        if not self.ws:
            raise RuntimeError("Not connected")

        pending_tasks: set[asyncio.Task] = set()

        try:
            async for raw_message in self.ws:
                # Rate limit incoming messages to prevent DoS
                if not self._incoming_limiter.allow():
                    continue

                try:
                    message = json.loads(raw_message)
                    logger.debug(f"Received: {message.get('method', 'response')}")

                    # Handle messages concurrently so long-running RPCs (like
                    # kernel.execute) don't block other requests (like kernel.inputReply)
                    task = asyncio.create_task(self._handle_and_respond(message))
                    pending_tasks.add(task)
                    task.add_done_callback(pending_tasks.discard)

                except json.JSONDecodeError:
                    logger.error(f"Invalid JSON received: {raw_message[:100]}")
                except Exception as e:
                    logger.error(f"Error handling message: {e}", exc_info=True)

            # Loop exited normally - check websocket close state
            # In websockets 15.x, the loop exits cleanly rather than raising ConnectionClosed
            close_code = self.ws.close_code
            close_reason = self.ws.close_reason or ""
            _log.info("ws.close", code=close_code, reason=close_reason)
            if close_code == 1000 and close_reason == "New connection from same runner":
                logger.info("Disconnected: Another instance connected from a different terminal.")
                self.running = False
                self._should_reconnect = False

        except ConnectionClosed as e:
            # Fallback for older websockets versions that raise an exception
            close_code = getattr(e.rcvd, "code", None) if e.rcvd else e.code
            close_reason = getattr(e.rcvd, "reason", None) if e.rcvd else e.reason
            _log.info("ws.close", code=close_code, reason=str(close_reason or ""))
            if close_code == 1000 and close_reason == "New connection from same runner":
                logger.info("Disconnected: Another instance connected from a different terminal.")
                self.running = False
                self._should_reconnect = False
        except ConnectionClosedError as e:
            _log.warning("ws.close_error", error=str(e))
        except Exception as e:
            logger.error(f"Receive loop error: {e}", exc_info=True)
        finally:
            self._output_buffer.set_connected(False)
            # Let in-flight handlers finish (they represent already-received
            # messages), then cancel any that take too long
            if pending_tasks:
                done, remaining = await asyncio.wait(
                    pending_tasks, timeout=2.0
                )
                for task in remaining:
                    task.cancel()
                if remaining:
                    await asyncio.gather(*remaining, return_exceptions=True)

    async def _get_http_client(self) -> httpx.AsyncClient:
        """Get or create the shared HTTP client for heartbeats."""
        if self._http_client is None:
            self._http_client = httpx.AsyncClient(timeout=10.0)
        return self._http_client

    async def _send_heartbeat(self):
        """Send a single heartbeat to the backend API."""
        client = await self._get_http_client()
        payload: dict = {"token": self.config.device_token}
        if self._metadata_supplier:
            try:
                metadata = self._metadata_supplier()
                if metadata:
                    payload["metadata"] = metadata
            except Exception as e:
                logger.debug(f"Metadata supplier error: {e}")
        response = await client.post(
            f"{self.config.api_url}/api/runner/heartbeat",
            json=payload,
        )
        if response.status_code != 200:
            logger.warning(f"Heartbeat failed: {response.status_code}")
        else:
            _log.debug("heartbeat.sent")
            # Cloud runner: server may signal shutdown when balance is exhausted
            try:
                data = response.json()
                if data.get("shouldShutdown"):
                    logger.warning("Server requested shutdown (balance exhausted)")
                    await self.message_handler.send_notification(
                        "runner.shutting_down",
                        {"reason": "balance_exhausted"},
                    )
                    await asyncio.sleep(1)
                    self.message_handler.request_shutdown("balance_exhausted")
            except Exception:
                pass

    async def heartbeat_loop(self):
        """Send periodic heartbeat to backend API."""
        while self.running:
            try:
                await self._send_heartbeat()
            except Exception as e:
                logger.error(f"Heartbeat error: {e}")
                # Reset client on error to get a fresh connection next time
                if self._http_client:
                    try:
                        await self._http_client.aclose()
                    except Exception:
                        pass
                    self._http_client = None

            await asyncio.sleep(10)  # Send heartbeat every 10 seconds

    async def run_with_reconnect(self):
        """Run the client with automatic reconnection."""
        self.running = True
        self._should_reconnect = True
        self._tasks: list[asyncio.Task] = []

        while self.running:
            try:
                connected = await self.connect()

                if connected:
                    # Notify caller of successful connection
                    if self.on_connect:
                        self.on_connect()

                    # Run heartbeat and receive loop concurrently as tasks
                    receive_task = asyncio.create_task(self.receive_loop())
                    heartbeat_task = asyncio.create_task(self.heartbeat_loop())
                    self._tasks = [receive_task, heartbeat_task]

                    # Wait for receive_loop to exit (connection closed), then cancel heartbeat
                    try:
                        await receive_task
                    except asyncio.CancelledError:
                        pass
                    finally:
                        # Cancel heartbeat when receive loop exits
                        heartbeat_task.cancel()
                        try:
                            await heartbeat_task
                        except asyncio.CancelledError:
                            pass

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Client error: {e}", exc_info=True)

            if self.running and self._should_reconnect:
                # Exponential backoff with jitter for reconnection
                jitter = random.uniform(0, self._reconnect_delay * 0.5)
                delay = self._reconnect_delay + jitter
                _log.info(
                    "ws.reconnect",
                    attempt=self._reconnect_delay,
                    delay_ms=round(delay * 1000),
                    url=self.config.ws_url,
                )
                try:
                    await asyncio.sleep(delay)
                except asyncio.CancelledError:
                    break
                self._reconnect_delay = min(
                    self._reconnect_delay * 2, self._max_reconnect_delay
                )
            elif not self._should_reconnect:
                logger.info("Runner shut down due to replacement.")
                break

