"""
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

logger = logging.getLogger(__name__)

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


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 = 60.0,
        heartbeat: float = 30.0,
    ) -> None:
        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()
        self._session = aiohttp.ClientSession()
        self._ws = await self._session.ws_connect(
            ws_url,
            ssl=ssl,
            headers=headers,
            heartbeat=self._heartbeat,
            receive_timeout=self._receive_timeout + 5.0,
        )
        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 = 60.0,
    heartbeat: float = 30.0,
) -> 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", "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", "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 = 60.0,
) -> 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", "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.

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

    Raises:
        asyncio.TimeoutError: If no terminal event within timeout.
        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()

    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=30.0,
                receive_timeout=timeout + 5.0,
            ) as ws:
                await ws.send_str(json.dumps({"action": "subscribe", "job_id": job_id}))
                deadline = asyncio.get_event_loop().time() + timeout
                while True:
                    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} within {timeout}s"
                        ) from None
                    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}, 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
