"""Production WebSocket client for AiCippy real-time streaming.

Connects to the WebSocket endpoint at wss://ws.aicippy.com,
handles authentication via JWT query parameter, automatic reconnection with
exponential backoff, ping/pong keepalive, and thread-safe message dispatch.
"""

from __future__ import annotations

import asyncio
import contextlib
import time
from collections.abc import AsyncIterator
from typing import Any, Final

import orjson
import structlog
import websockets
import websockets.asyncio.client
from websockets.asyncio.client import ClientConnection
from websockets.exceptions import (
    ConnectionClosed,
    ConnectionClosedError,
    ConnectionClosedOK,
    InvalidStatus,
)

from aicippy.websocket.models import (
    ChatRequest,
    ConnectionState,
    StreamToken,
    WebSocketMessage,
)

logger: structlog.stdlib.BoundLogger = structlog.get_logger(__name__)


def _safe_int(value: Any) -> int:
    """Convert a value to int, returning 0 on failure."""
    try:
        return int(value)
    except (ValueError, TypeError):
        return 0


# ============================================================================
# Constants
# ============================================================================

DEFAULT_WS_URL: Final[str] = "wss://ws.aicippy.com"
MAX_RECONNECT_ATTEMPTS: Final[int] = 5
INITIAL_BACKOFF_SECONDS: Final[float] = 1.0
MAX_BACKOFF_SECONDS: Final[float] = 30.0
BACKOFF_MULTIPLIER: Final[float] = 2.0
PING_INTERVAL_SECONDS: Final[float] = 30.0
PING_TIMEOUT_SECONDS: Final[float] = 10.0
RECEIVE_TIMEOUT_SECONDS: Final[float] = 60.0
MAX_MESSAGE_SIZE: Final[int] = 10 * 1024 * 1024

# ============================================================================
# Exceptions
# ============================================================================


class WebSocketNotConnectedError(ConnectionError):
    """Raised when the WebSocket is not connected."""

    def __init__(self) -> None:
        """Initialize exception with default message."""
        super().__init__("WebSocket is not connected")


class NoResponseQueueError(ConnectionError):
    """Raised when no response queue exists for a correlation ID."""

    def __init__(self, correlation_id: str) -> None:
        """Initialize exception with formatted message."""
        super().__init__(f"No response queue for correlation_id={correlation_id}. Call send_chat first.")


class ConnectionAttemptsExhaustedError(ConnectionError):
    """Raised when max reconnection attempts are reached."""

    def __init__(self, attempts: int, last_error: Exception | None) -> None:
        """Initialize exception with formatted message."""
        super().__init__(f"Failed to connect after {attempts} attempts: {last_error}")


class WebSocketUnavailableError(ConnectionError):
    """Raised when the WebSocket connection is not available."""

    def __init__(self) -> None:
        """Initialize exception with default message."""
        super().__init__("WebSocket connection is not available")


# ============================================================================
# WebSocket Client
# ============================================================================


class WebSocketClient:
    """Production-grade WebSocket client for AiCippy streaming.

    Features:
        - JWT-authenticated connection via query parameter
        - Automatic reconnection with exponential backoff (max 5 attempts)
        - Ping/pong keepalive every 30 seconds
        - Thread-safe message sending via asyncio
        - AsyncIterator-based token streaming per correlation_id
        - Structured logging of all connection lifecycle events

    Usage:
        >>> client = WebSocketClient()
        >>> await client.connect(auth_token="eyJ...")
        >>> await client.send_chat("Hello", model="odin", ...)
        >>> async for token in client.stream_response(correlation_id):
        ...     print(token.text, end="")
        >>> await client.disconnect()
    """

    __slots__ = (
        "_auth_token",
        "_connection",
        "_keepalive_task",
        "_listener_task",
        "_lock",
        "_reconnect_count",
        "_response_queues",
        "_session_id",
        "_state",
        "_url",
    )

    def __init__(
        self,
        url: str = DEFAULT_WS_URL,
        auth_token: str = "",
        session_id: str = "",
    ) -> None:
        """Initialize WebSocket client.

        Args:
            url: WebSocket endpoint URL. Defaults to wss://ws.aicippy.com/prod.
            auth_token: Bearer token from Cognito SSO for authentication.
            session_id: Session identifier for message correlation.

        """
        self._url: str = url
        self._auth_token: str = auth_token
        self._session_id: str = session_id
        self._connection: ClientConnection | None = None
        self._state: ConnectionState = ConnectionState.DISCONNECTED
        self._lock: asyncio.Lock | None = None
        self._reconnect_count: int = 0
        self._response_queues: dict[str, asyncio.Queue[StreamToken | None]] = {}
        self._listener_task: asyncio.Task[None] | None = None
        self._keepalive_task: asyncio.Task[None] | None = None

    # ========================================================================
    # Properties
    # ========================================================================

    @property
    def connected(self) -> bool:
        """Whether the WebSocket connection is currently established."""
        return self._state == ConnectionState.CONNECTED and self._connection is not None

    @property
    def state(self) -> ConnectionState:
        """Current connection state."""
        return self._state

    # ========================================================================
    # Connection Lifecycle
    # ========================================================================

    async def connect(self, auth_token: str) -> None:
        """Establish an authenticated WebSocket connection.

        Connects to the configured endpoint with the JWT token passed as a
        query parameter. Starts the background listener and keepalive tasks.

        Args:
            auth_token: Valid JWT token for authentication.

        Raises:
            ConnectionError: If the connection cannot be established after retries.

        """
        if self.connected:
            logger.debug("websocket_already_connected")
            return

        self._auth_token = auth_token
        self._state = ConnectionState.CONNECTING
        logger.info("websocket_connecting", url=self._url)

        await self._establish_connection()

        self._listener_task = asyncio.create_task(
            self._listen_loop(),
            name="ws-listener",
        )
        self._keepalive_task = asyncio.create_task(
            self._keepalive_loop(),
            name="ws-keepalive",
        )

        logger.info("websocket_connected", url=self._url)

    async def disconnect(self) -> None:
        """Disconnect the WebSocket connection cleanly.

        Cancels background tasks, closes the connection, and drains all
        pending response queues with a sentinel value.
        """
        logger.info("websocket_disconnecting")
        self._state = ConnectionState.DISCONNECTED

        # Cancel background tasks
        for task in (self._listener_task, self._keepalive_task):
            if task is not None and not task.done():
                task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await task

        self._listener_task = None
        self._keepalive_task = None

        # Close the WebSocket connection
        if self._connection is not None:
            try:
                await self._connection.close()
            except (OSError, websockets.exceptions.WebSocketException):
                logger.debug("websocket_close_error", exc_info=True)
            finally:
                self._connection = None

        # Signal all waiting consumers that the connection is gone
        for queue in self._response_queues.values():
            with contextlib.suppress(asyncio.QueueFull):
                queue.put_nowait(None)

        self._response_queues.clear()
        logger.info("websocket_disconnected")

    # ========================================================================
    # Messaging
    # ========================================================================

    async def send_chat(
        self,
        message: str,
        model: str,
        system_prompt: str,
        conversation_history: list[dict[str, Any]],
        correlation_id: str,
    ) -> None:
        """Send a chat message over the WebSocket.

        Creates a response queue for the given correlation_id before sending,
        so that stream_response can immediately begin consuming tokens.

        Args:
            message: The user's message text.
            model: Model identifier.
            system_prompt: System prompt for the model.
            conversation_history: Prior conversation turns.
            correlation_id: Unique ID to match response tokens.

        Raises:
            ConnectionError: If not connected.

        """
        if not self.connected:
            raise ConnectionError("WebSocket is not connected")  # noqa: TRY003

        # Create the response queue before sending to avoid race conditions
        self._response_queues[correlation_id] = asyncio.Queue()

        request = ChatRequest(
            message=message,
            model=model,
            system_prompt=system_prompt,
            conversation_history=conversation_history,
            correlation_id=correlation_id,
        )

        await self._send_raw(request.to_wire())
        logger.debug(
            "websocket_chat_sent",
            correlation_id=correlation_id,
            model=model,
        )

    async def stream_response(self, correlation_id: str) -> AsyncIterator[StreamToken]:
        """Yield stream tokens for a given correlation_id as they arrive.

        Blocks until each token is available. Yields until the final token
        (done=True) is received or the connection is lost (sentinel None).

        Args:
            correlation_id: The correlation ID from send_chat.

        Yields:
            StreamToken instances as they arrive from the server.

        Raises:
            ConnectionError: If no queue exists for the correlation_id.

        """
        queue = self._response_queues.get(correlation_id)
        if queue is None:
            raise ConnectionError(  # noqa: TRY003
                f"No response queue for correlation_id={correlation_id}. Call send_chat first.",
            )

        try:
            while True:
                token = await queue.get()
                if token is None:
                    # Sentinel: connection lost or stream cancelled
                    break
                yield token
                if token.done:
                    break
        finally:
            # Clean up the queue once consumed
            self._response_queues.pop(correlation_id, None)

    async def ping(self) -> float:
        """Send a WebSocket ping and measure round-trip latency.

        Returns:
            Latency in milliseconds.

        Raises:
            ConnectionError: If not connected.

        """
        if not self.connected or self._connection is None:
            raise ConnectionError("WebSocket is not connected")  # noqa: TRY003

        start = time.monotonic()
        pong_waiter = await self._connection.ping()
        await asyncio.wait_for(pong_waiter, timeout=PING_TIMEOUT_SECONDS)
        elapsed_ms = (time.monotonic() - start) * 1000.0

        logger.debug("websocket_ping", latency_ms=round(elapsed_ms, 2))
        return elapsed_ms

    # ========================================================================
    # Internal: Connection Management
    # ========================================================================

    async def _establish_connection(self) -> None:
        """Establish the raw WebSocket connection with retry logic.

        Uses exponential backoff up to MAX_RECONNECT_ATTEMPTS. On exhaustion,
        raises ConnectionError.
        """
        backoff = INITIAL_BACKOFF_SECONDS
        last_error: Exception | None = None

        for attempt in range(1, MAX_RECONNECT_ATTEMPTS + 1):
            try:
                ws_url = f"{self._url}?token={self._auth_token}"
                self._connection = await websockets.asyncio.client.connect(
                    ws_url,
                    ping_interval=None,  # We manage ping/pong ourselves
                    ping_timeout=None,
                    close_timeout=10,
                    max_size=10 * 1024 * 1024,  # 10 MiB max message
                )
                self._state = ConnectionState.CONNECTED
                self._reconnect_count = 0
                return  # noqa: TRY300

            except (OSError, InvalidStatus) as exc:
                last_error = exc
                logger.warning(
                    "websocket_connect_failed",
                    attempt=attempt,
                    max_attempts=MAX_RECONNECT_ATTEMPTS,
                    backoff_seconds=backoff,
                    error=str(exc),
                )
                if attempt < MAX_RECONNECT_ATTEMPTS:
                    await asyncio.sleep(backoff)
                    backoff = min(backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF_SECONDS)

        self._state = ConnectionState.DISCONNECTED
        raise ConnectionError(  # noqa: TRY003
            f"Failed to connect after {MAX_RECONNECT_ATTEMPTS} attempts: {last_error}",
        )

    async def _reconnect(self) -> None:
        """Attempt to re-establish the WebSocket connection transparently.

        Called by the listener loop when the connection drops unexpectedly.
        """
        if self._state == ConnectionState.DISCONNECTED:
            # Deliberate disconnect; do not reconnect
            return

        self._state = ConnectionState.RECONNECTING
        self._reconnect_count += 1
        logger.info(
            "websocket_reconnecting",
            attempt=self._reconnect_count,
        )

        try:
            await self._establish_connection()
            logger.info("websocket_reconnected")
        except ConnectionError:
            logger.exception("websocket_reconnect_exhausted")
            self._state = ConnectionState.DISCONNECTED
            # Signal all waiting consumers
            for queue in self._response_queues.values():
                with contextlib.suppress(asyncio.QueueFull):
                    queue.put_nowait(None)

    # ========================================================================
    # Internal: Background Loops
    # ========================================================================

    async def _listen_loop(self) -> None:
        """Background task that receives WebSocket messages and dispatches them.

        Automatically triggers reconnection on unexpected connection closure.
        """
        while self._state in (ConnectionState.CONNECTED, ConnectionState.RECONNECTING):
            try:
                if self._connection is None:
                    await asyncio.sleep(0.1)
                    continue

                raw = await self._connection.recv()
                if isinstance(raw, bytes):
                    data: dict[str, Any] = orjson.loads(raw)
                else:
                    data = orjson.loads(raw.encode("utf-8"))

                await self._dispatch_message(data)

            except ConnectionClosedOK:
                logger.info("websocket_closed_ok")
                break

            except (ConnectionClosed, ConnectionClosedError) as exc:
                logger.warning("websocket_connection_lost", error=str(exc))
                self._connection = None
                await self._reconnect()
                if self._state == ConnectionState.DISCONNECTED:
                    break

            except asyncio.CancelledError:
                raise

            except (ValueError, TypeError):
                logger.error("websocket_listen_error", exc_info=True)
                await asyncio.sleep(0.5)

    async def _keepalive_loop(self) -> None:
        """Background task that sends periodic pings to keep the connection alive.

        API Gateway WebSocket connections idle-timeout after ~10 minutes.
        Pinging every 30 seconds prevents this.
        """
        while self._state in (ConnectionState.CONNECTED, ConnectionState.RECONNECTING):
            try:
                await asyncio.sleep(PING_INTERVAL_SECONDS)
                if self.connected and self._connection is not None:
                    pong_waiter = await self._connection.ping()
                    await asyncio.wait_for(
                        pong_waiter,
                        timeout=PING_TIMEOUT_SECONDS,
                    )
                    logger.debug("websocket_keepalive_pong")

            except asyncio.CancelledError:
                raise

            except TimeoutError:
                logger.warning("websocket_keepalive_timeout")
                # Force reconnect on ping timeout
                if self._connection is not None:
                    with contextlib.suppress(OSError, websockets.exceptions.WebSocketException):
                        await self._connection.close()
                    self._connection = None

            except (ConnectionClosed, ConnectionClosedError):
                logger.warning("websocket_keepalive_connection_lost")
                break

            except (OSError, RuntimeError):
                logger.error("websocket_keepalive_error", exc_info=True)

    # ========================================================================
    # Internal: Message Dispatch
    # ========================================================================

    async def _dispatch_message(self, data: dict[str, Any]) -> None:
        """Route an incoming WebSocket message to the appropriate handler.

        Args:
            data: Parsed JSON message from the server.

        """
        msg = WebSocketMessage.from_dict(data)

        if msg.type == "token":
            await self._handle_token(msg)
        elif msg.type == "error":
            await self._handle_error(msg)
        else:
            logger.debug(
                "websocket_unknown_message_type",
                type=msg.type,
                correlation_id=msg.correlation_id,
            )

    async def _handle_token(self, msg: WebSocketMessage) -> None:
        """Handle an incoming stream token message.

        Converts the payload to a StreamToken and enqueues it for the
        matching correlation_id consumer.
        """
        correlation_id = msg.correlation_id
        queue = self._response_queues.get(correlation_id)
        if queue is None:
            logger.warning(
                "websocket_orphan_token",
                correlation_id=correlation_id,
            )
            return

        payload = msg.payload
        done = bool(payload.get("done", False))

        token = StreamToken(
            text=str(payload.get("text", "")),
            done=done,
            input_tokens=_safe_int(payload.get("input_tokens", 0)) if done else 0,
            output_tokens=_safe_int(payload.get("output_tokens", 0)) if done else 0,
        )

        await queue.put(token)

    async def _handle_error(self, msg: WebSocketMessage) -> None:
        """Handle an incoming error message from the server.

        Logs the error and signals the consumer queue with a sentinel.
        """
        correlation_id = msg.correlation_id
        error_text = msg.payload.get("message", "Unknown server error")

        logger.error(
            "websocket_server_error",
            correlation_id=correlation_id,
            error=error_text,
        )

        queue = self._response_queues.get(correlation_id)
        if queue is not None:
            with contextlib.suppress(asyncio.QueueFull):
                queue.put_nowait(None)

    # ========================================================================
    # Internal: Raw Send
    # ========================================================================

    async def _send_raw(self, data: dict[str, Any]) -> None:
        """Serialize and send a message over the WebSocket.

        Thread-safe via asyncio lock.

        Args:
            data: JSON-serializable dictionary to send.

        Raises:
            ConnectionError: If not connected.

        """
        if self._lock is None:
            self._lock = asyncio.Lock()
        async with self._lock:
            if self._connection is None:
                raise ConnectionError("WebSocket connection is not available")  # noqa: TRY003
            payload = orjson.dumps(data)
            await self._connection.send(payload)
