"""
WebSocket client for job completion push and bidirectional channel (e.g. /ws).

Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com

Use when the server exposes a /ws endpoint. Supports:
- wait_for_job_via_websocket(): one-shot subscribe and wait for terminal event.
- open_bidirectional_ws_channel(): bidirectional channel (send + receive loop).
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, AsyncIterator, Dict, Optional

from mcp_proxy_adapter.client.jsonrpc_client.transport import JsonRpcTransport
from mcp_proxy_adapter.core.constants import (
    WS_HEARTBEAT_INTERVAL,
    WS_RECEIVE_TIMEOUT,
    WS_RECEIVE_TIMEOUT_BUFFER,
    resolve_timeout,
)

logger = logging.getLogger(__name__)

# Events that indicate terminal job state (stop waiting)
WS_TERMINAL_EVENTS = frozenset({"job_completed", "job_failed", "job_stopped"})

# Canonical terminal delivery events (ws_delivery subscription, keyed by ws_id)
WS_DELIVERY_TERMINAL_EVENTS = frozenset(
    {"delivery_completed", "delivery_failed", "delivery_cancelled"}
)


class BidirectionalWsChannel:
    """
    Bidirectional WebSocket channel: send and receive JSON messages over /ws.

    Use as async context manager. Same auth (headers, TLS) as transport.
    """

    def __init__(
        self,
        transport: JsonRpcTransport,
        receive_timeout: float = WS_RECEIVE_TIMEOUT,
        heartbeat: float = WS_HEARTBEAT_INTERVAL,
    ) -> None:
        """
        Initialize bidirectional WebSocket channel.

        Args:
            transport: JSON-RPC transport for connection details.
            receive_timeout: Timeout for receiving messages in seconds.
            heartbeat: Heartbeat interval in seconds.
        """
        self._transport = transport
        self._receive_timeout = receive_timeout
        self._heartbeat = heartbeat
        self._ws: Any = None
        self._session: Any = None

    async def __aenter__(self) -> "BidirectionalWsChannel":
        import aiohttp

        ws_url = self._transport.ws_url
        headers = {
            k: v
            for k, v in self._transport.headers.items()
            if k.lower() != "content-type"
        }
        ssl = self._transport.get_ws_ssl_context()
        hb = resolve_timeout(self._heartbeat)
        rt = resolve_timeout(self._receive_timeout)
        if rt is not None:
            rt += WS_RECEIVE_TIMEOUT_BUFFER
        self._session = aiohttp.ClientSession()
        self._ws = await self._session.ws_connect(
            ws_url,
            ssl=ssl,
            headers=headers,
            heartbeat=hb,
            receive_timeout=rt,
        )
        return self

    async def __aexit__(self, *args: Any) -> None:
        if self._ws:
            await self._ws.close()
            self._ws = None
        if self._session:
            await self._session.close()
            self._session = None

    async def send_json(self, payload: Dict[str, Any]) -> None:
        """Send a JSON object to the server."""
        if self._ws is None:
            raise RuntimeError("Channel not connected; use async with first")
        await self._ws.send_str(json.dumps(payload))

    async def receive_iter(self) -> AsyncIterator[Dict[str, Any]]:
        """
        Async iterator over incoming JSON messages (server -> client).

        Yields each parsed JSON object. Stops when the connection closes or errors.
        """
        import aiohttp

        if self._ws is None:
            raise RuntimeError("Channel not connected; use async with first")
        while True:
            msg = await self._ws.receive()
            if msg.type == aiohttp.WSMsgType.TEXT:
                try:
                    yield json.loads(msg.data)
                except json.JSONDecodeError:
                    logger.warning("WS non-JSON: %s", msg.data[:200])
            elif msg.type in (
                aiohttp.WSMsgType.CLOSE,
                aiohttp.WSMsgType.ERROR,
            ):
                break


def open_bidirectional_ws_channel(
    transport: JsonRpcTransport,
    receive_timeout: float = WS_RECEIVE_TIMEOUT,
    heartbeat: float = WS_HEARTBEAT_INTERVAL,
) -> BidirectionalWsChannel:
    """
    Create a bidirectional WebSocket channel (async context manager) for /ws.

    Returns a BidirectionalWsChannel. Use with async with; then call
    send_json() to send and receive_iter() to consume server messages.

    Example:
        channel = open_bidirectional_ws_channel(client)
        async with channel:
            await channel.send_json({
                "action": "subscribe",
                "subscription_type": "queue_job",
                "job_id": job_id,
            })
            async for msg in channel.receive_iter():
                print("Received", msg)
                if msg.get("event") in ("job_completed", "job_failed"):
                    break
            await channel.send_json({
                "action": "unsubscribe",
                "subscription_type": "queue_job",
                "job_id": job_id,
            })
    """
    return BidirectionalWsChannel(
        transport=transport,
        receive_timeout=receive_timeout,
        heartbeat=heartbeat,
    )


async def wait_for_job_via_websocket(
    transport: JsonRpcTransport,
    job_id: str,
    timeout: float = WS_RECEIVE_TIMEOUT,
    heartbeat: float = WS_HEARTBEAT_INTERVAL,
) -> Dict[str, Any]:
    """
    Connect to /ws, subscribe to job_id, and wait for job_completed / job_failed.

    Uses the same base URL, auth headers, and TLS as the HTTP transport.
    Server is expected to accept WebSocket at /ws and support:
    - Outbound: {"action": "subscribe", "subscription_type": "queue_job", "job_id": "<job_id>"}
    - Inbound: {"event": "job_completed"|"job_failed"|..., "job_id": "...", ...}

    Args:
        transport: JsonRpcTransport (or JsonRpcClient) with base_url, headers, SSL.
        job_id: Job identifier to subscribe to.
        timeout: Max seconds to wait for a terminal event. 0 = no timeout.
        heartbeat: WebSocket heartbeat interval in seconds. 0 = disabled.

    Returns:
        Last message payload (e.g. event, job_id, result or error).

    Raises:
        asyncio.TimeoutError: If no terminal event within timeout (when timeout > 0).
        RuntimeError: On connection or protocol errors.
    """
    import aiohttp

    ws_url = transport.ws_url
    headers = {
        k: v for k, v in transport.headers.items() if k.lower() != "content-type"
    }
    ssl = transport.get_ws_ssl_context()

    effective_timeout = resolve_timeout(timeout)
    effective_hb = resolve_timeout(heartbeat)

    ws_receive_timeout: Optional[float] = None
    if effective_timeout is not None:
        ws_receive_timeout = effective_timeout + WS_RECEIVE_TIMEOUT_BUFFER

    result: Optional[Dict[str, Any]] = None

    async with aiohttp.ClientSession() as session:
        try:
            async with session.ws_connect(
                ws_url,
                ssl=ssl,
                headers=headers,
                heartbeat=effective_hb,
                receive_timeout=ws_receive_timeout,
            ) as ws:
                await ws.send_str(
                    json.dumps(
                        {
                            "action": "subscribe",
                            "subscription_type": "queue_job",
                            "job_id": job_id,
                        }
                    )
                )

                deadline: Optional[float] = None
                if effective_timeout is not None:
                    deadline = asyncio.get_event_loop().time() + effective_timeout

                while True:
                    if deadline is not None:
                        remaining = max(0.1, deadline - asyncio.get_event_loop().time())
                        try:
                            msg = await asyncio.wait_for(
                                ws.receive(), timeout=remaining
                            )
                        except asyncio.TimeoutError:
                            raise asyncio.TimeoutError(
                                f"No job completion event for job_id={job_id} "
                                f"within {timeout}s"
                            ) from None
                    else:
                        msg = await ws.receive()

                    if msg.type == aiohttp.WSMsgType.TEXT:
                        try:
                            data = json.loads(msg.data)
                        except json.JSONDecodeError:
                            logger.warning("WS non-JSON: %s", msg.data[:200])
                            continue
                        result = data
                        ev = data.get("event") or data.get("type")
                        if ev in WS_TERMINAL_EVENTS:
                            logger.debug("WS terminal event: %s", ev)
                            return result
                        logger.debug("WS event: %s", ev)
                    elif msg.type in (
                        aiohttp.WSMsgType.CLOSE,
                        aiohttp.WSMsgType.ERROR,
                    ):
                        raise RuntimeError(
                            f"WebSocket closed or error: type={msg.type}, "
                            f"data={getattr(msg, 'data', None)}"
                        )
        except aiohttp.ClientError as e:
            raise RuntimeError(f"WebSocket connection failed: {e}") from e

    if result is None:
        raise asyncio.TimeoutError(
            f"No job completion event for job_id={job_id} within {timeout}s"
        )
    return result


async def wait_for_terminal_delivery_via_websocket(
    transport: JsonRpcTransport,
    ws_id: str,
    *,
    timeout: float = WS_RECEIVE_TIMEOUT,
    heartbeat: float = WS_HEARTBEAT_INTERVAL,
) -> Dict[str, Any]:
    """
    Connect to /ws, subscribe by ws_id (subscription_type ws_delivery), wait for terminal.

    Subscribe-after-acceptance: call only after server returned ws_id in acceptance.
    Server may replay terminal immediately if execution already finished. Returns
    payload matching WsDeliveryTerminalPayload shape (event, ws_id, optional job_id,
    result, error).

    Args:
        transport: JsonRpcTransport with base_url, headers, SSL.
        ws_id: WebSocket delivery identity from acceptance (server-issued).
        timeout: Max seconds to wait for terminal delivery. 0 = no timeout.
        heartbeat: WebSocket heartbeat interval. 0 = disabled.

    Returns:
        Dict with event, ws_id, optional job_id, result, error.

    Raises:
        asyncio.TimeoutError: No terminal event within timeout.
        RuntimeError: Connection or protocol error; or terminal for different ws_id.
    """
    import aiohttp

    ws_url = transport.ws_url
    headers = {
        k: v for k, v in transport.headers.items() if k.lower() != "content-type"
    }
    ssl = transport.get_ws_ssl_context()
    effective_timeout = resolve_timeout(timeout)
    effective_hb = resolve_timeout(heartbeat)
    ws_receive_timeout: Optional[float] = None
    if effective_timeout is not None:
        ws_receive_timeout = effective_timeout + WS_RECEIVE_TIMEOUT_BUFFER

    async with aiohttp.ClientSession() as session:
        try:
            async with session.ws_connect(
                ws_url,
                ssl=ssl,
                headers=headers,
                heartbeat=effective_hb,
                receive_timeout=ws_receive_timeout,
            ) as ws:
                await ws.send_str(
                    json.dumps(
                        {
                            "action": "subscribe",
                            "subscription_type": "ws_delivery",
                            "ws_id": ws_id,
                        }
                    )
                )
                deadline: Optional[float] = None
                if effective_timeout is not None:
                    deadline = asyncio.get_event_loop().time() + effective_timeout

                while True:
                    if deadline is not None:
                        remaining = max(0.1, deadline - asyncio.get_event_loop().time())
                        try:
                            msg = await asyncio.wait_for(
                                ws.receive(), timeout=remaining
                            )
                        except asyncio.TimeoutError:
                            raise asyncio.TimeoutError(
                                f"No terminal delivery for ws_id={ws_id} "
                                f"within {timeout}s"
                            ) from None
                    else:
                        msg = await ws.receive()

                    if msg.type == aiohttp.WSMsgType.TEXT:
                        try:
                            data = json.loads(msg.data)
                        except json.JSONDecodeError:
                            logger.warning(
                                "WS non-JSON: %s", getattr(msg, "data", "")[:200]
                            )
                            continue
                        ev = data.get("event") or data.get("type")
                        if ev in WS_DELIVERY_TERMINAL_EVENTS:
                            payload_ws_id = data.get("ws_id") or data.get("id")
                            if payload_ws_id != ws_id:
                                raise RuntimeError(
                                    f"Terminal event ws_id {payload_ws_id!r} "
                                    f"does not match subscribed ws_id {ws_id!r}"
                                )
                            logger.debug("WS delivery terminal event: %s", ev)
                            return data
                    elif msg.type in (
                        aiohttp.WSMsgType.CLOSE,
                        aiohttp.WSMsgType.ERROR,
                    ):
                        raise RuntimeError(
                            f"WebSocket closed or error: type={msg.type}, "
                            f"data={getattr(msg, 'data', None)}"
                        )
        except aiohttp.ClientError as e:
            raise RuntimeError(f"WebSocket connection failed: {e}") from e

    raise asyncio.TimeoutError(
        f"No terminal delivery for ws_id={ws_id} within {timeout}s"
    )
