"""WebSocket connection lifecycle manager.

Manages the persistent WSS connection between qRemoteX and the qRunX
Executor Gateway.  Responsibilities:
  - Connect with API-key auth headers
  - Send/receive QRP messages
  - Heartbeat loop
  - Outbox (single-writer) pattern
  - Auto-reconnect with exponential backoff
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import signal
import time
from typing import Any, Callable, Coroutine

import aiohttp

from qremotex.config import ExecutorConfig
from qremotex.connection.auth import AuthManager
from qremotex.connection.mtls import MTLSConfig, build_ssl_context
from qremotex.connection.reconnect import ReconnectStrategy
from qremotex.plan_gating import handle_gateway_rejection, is_permanent_rejection
from qremotex.protocol import serializer
from qremotex.protocol.messages import (
    MessageType,
    heartbeat_message,
    hello_message,
    pong_message,
)

logger = logging.getLogger("qremotex.connection")

# Type alias for the message handler callback
MessageHandler = Callable[[dict[str, Any]], Coroutine[Any, Any, None]]


class ConnectionManager:
    """Persistent WebSocket connection to the qRunX Executor Gateway."""

    def __init__(
        self,
        config: ExecutorConfig,
        on_message: MessageHandler,
        metrics_fn: Callable[[], dict[str, Any]] | None = None,
    ) -> None:
        self._config = config
        self._auth = AuthManager(config)
        self._reconnect = ReconnectStrategy(max_sec=config.reconnect_max_sec)
        self._on_message = on_message
        self._metrics_fn = metrics_fn

        # mTLS configuration (Enterprise)
        self._mtls_config = MTLSConfig()
        self._ssl_context = build_ssl_context(self._mtls_config)

        # Internal state
        self._stop = asyncio.Event()
        self._draining = False
        self._outbox: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        self._send_lock = asyncio.Lock()
        self._ws: aiohttp.ClientWebSocketResponse | None = None
        self._connected = asyncio.Event()

    # ------------------------------------------------------------------ #
    # Public API
    # ------------------------------------------------------------------ #
    async def run(self) -> None:
        """Main event loop — connect, serve, reconnect forever."""
        loop = asyncio.get_running_loop()
        for sig in (signal.SIGINT, signal.SIGTERM):
            try:
                loop.add_signal_handler(sig, self._stop.set)
            except NotImplementedError:
                pass  # Windows

        while not self._stop.is_set():
            try:
                async with aiohttp.ClientSession() as session:
                    await self._serve(session)
                self._reconnect.reset()
            except asyncio.CancelledError:
                break
            except Exception as exc:
                # Check for plan gating rejection
                close_code = getattr(exc, "close_code", None)
                close_reason = str(exc) if close_code else ""
                plan_error = handle_gateway_rejection(close_code, close_reason)
                if plan_error and is_permanent_rejection(plan_error):
                    logger.error(
                        "Permanent rejection from gateway: %s. Not reconnecting.",
                        plan_error.message,
                    )
                    break

                delay = self._reconnect.next_delay()
                logger.warning(
                    "Connection lost (%s: %s). Reconnecting in %.1fs (attempt %d)",
                    type(exc).__name__,
                    exc,
                    delay,
                    self._reconnect.attempt,
                )
                try:
                    await asyncio.wait_for(self._stop.wait(), timeout=delay)
                    break  # stop was set during wait
                except asyncio.TimeoutError:
                    pass  # timeout expired, retry

        logger.info("Connection manager shutting down")

    async def send(self, msg: dict[str, Any]) -> None:
        """Enqueue an outbound QRP message."""
        # Attach HMAC signature
        self._auth.sign_dict(msg)
        await self._outbox.put(msg)

    async def stop(self) -> None:
        self._stop.set()

    @property
    def is_connected(self) -> bool:
        return self._connected.is_set()

    @property
    def is_draining(self) -> bool:
        return self._draining

    def start_drain(self) -> None:
        """Enter drain mode — no new tasks will be accepted."""
        self._draining = True
        logger.info("Drain mode activated — no new tasks will be accepted")

    # ------------------------------------------------------------------ #
    # Connection lifecycle
    # ------------------------------------------------------------------ #
    async def _serve(self, session: aiohttp.ClientSession) -> None:
        """Single connection lifecycle: connect → serve → cleanup."""
        ws = await self._ws_connect(session)
        self._ws = ws
        self._connected.set()
        self._reconnect.reset()
        logger.info("Connected to %s", self._config.gateway_url)

        # Send hello
        hello = hello_message(
            executor_id=self._config.executor_id,
            api_key_id=self._auth.key_id,
            version=self._config.version,
            capabilities=self._config.build_capabilities(),
            labels=self._config.labels,
            system_info=self._config.build_system_info(),
            name=self._config.executor_name,
            max_concurrency=self._config.concurrency,
            group=self._config.group,
        )
        await self.send(hello)

        # Launch background tasks
        sender_task = asyncio.create_task(self._sender_loop(ws))
        hb_task = asyncio.create_task(self._heartbeat_loop())

        try:
            async for msg in ws:
                if self._stop.is_set():
                    break
                if msg.type == aiohttp.WSMsgType.TEXT:
                    logger.debug(
                        "ws_message_received: type=%s data_len=%d",
                        msg.type, len(msg.data) if msg.data else 0,
                    )
                    parsed = serializer.deserialize(msg.data)
                    if parsed:
                        await self._handle_incoming(parsed)
                elif msg.type == aiohttp.WSMsgType.ERROR:
                    err = ws.exception()
                    logger.error("WebSocket error message: %s", err)
                    raise ConnectionError(f"WebSocket error: {err}")
                else:
                    # NOTE: CLOSE/CLOSING/CLOSED are intercepted by aiohttp's
                    # __anext__ (raises StopAsyncIteration) and never reach here.
                    logger.warning("ws_unexpected_msg_type: %s", msg.type)

            # async for exited — connection was closed by server or stop was set.
            # aiohttp intercepts CLOSE/CLOSING/CLOSED frames internally,
            # so the iterator simply exhausts when the connection drops.
            logger.warning(
                "ws_loop_exited: stop_set=%s ws_closed=%s close_code=%s",
                self._stop.is_set(), ws.closed, ws.close_code,
            )
            if not self._stop.is_set():
                raise ConnectionError(
                    f"WebSocket closed by server (close_code={ws.close_code})"
                )
        finally:
            self._connected.clear()
            self._ws = None
            for t in (sender_task, hb_task):
                t.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await t

    async def _ws_connect(self, session: aiohttp.ClientSession) -> aiohttp.ClientWebSocketResponse:
        """Establish the WebSocket connection with auth."""
        url = self._config.gateway_url
        params = self._auth.handshake_params()
        headers = self._auth.handshake_headers()

        # Build URL with query params
        if "?" in url:
            url += "&"
        else:
            url += "?"
        url += "&".join(f"{k}={v}" for k, v in params.items())

        return await session.ws_connect(
            url,
            # NOTE: Do NOT set heartbeat= here. aiohttp's WS-level PING/PONG
            # control frames are often blocked by infrastructure proxies (ALB,
            # CloudFront, nginx), causing spurious 1006 disconnects.
            # qRemoteX has its own application-level heartbeat via
            # _heartbeat_loop() which sends QRP heartbeat TEXT messages that
            # pass through any proxy reliably.
            heartbeat=None,
            headers=headers,
            ssl=self._ssl_context,
            # Allow large messages (e.g. execute_script with base64 PDFs
            # in input_vars).  Default is 4 MB which is too small for
            # payloads with embedded file data.
            max_msg_size=16 * 1024 * 1024,  # 16 MB
        )

    # ------------------------------------------------------------------ #
    # Background loops
    # ------------------------------------------------------------------ #
    async def _sender_loop(self, ws: aiohttp.ClientWebSocketResponse) -> None:
        """Drain outbox and send messages over the websocket."""
        while True:
            msg = await self._outbox.get()
            try:
                async with self._send_lock:
                    raw = serializer.serialize(msg)
                    await ws.send_str(raw)
            except Exception:
                logger.exception("Failed to send message")
                return

    async def _heartbeat_loop(self) -> None:
        """Send periodic heartbeat with system metrics."""
        while True:
            await asyncio.sleep(self._config.heartbeat_sec)
            metrics = self._metrics_fn() if self._metrics_fn else {}
            hb = heartbeat_message(
                executor_id=self._config.executor_id,
                ts=int(time.time()),
                metrics=metrics,
            )
            if self._draining:
                hb["draining"] = True
            await self.send(hb)

    # ------------------------------------------------------------------ #
    # Incoming message handling
    # ------------------------------------------------------------------ #
    async def _handle_incoming(self, data: dict[str, Any]) -> None:
        """Route incoming messages — handle protocol-level ones, delegate the rest."""
        msg_type = data.get("type")

        if msg_type == MessageType.PING.value:
            pong = pong_message(self._config.executor_id)
            await self.send(pong)
            return

        if msg_type == "hello_ack":
            # Server is the source of truth for executor_id (bound to API key).
            # Sync it locally so heartbeats/pongs use the correct ID.
            server_eid = data.get("executor_id")
            if server_eid and server_eid != self._config.executor_id:
                logger.info(
                    "Server assigned executor_id=%s (was %s)",
                    server_eid, self._config.executor_id,
                )
                self._config.executor_id = server_eid
            return

        # Delegate all other message types (execute_script, execute_tool, cancel, etc.)
        try:
            await self._on_message(data)
        except Exception:
            logger.exception("Error in message handler for type=%s", msg_type)
