#!/usr/bin/env python3
"""
Tests for WebSocket job status client (wait_for_job_via_websocket, BidirectionalWsChannel).

Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com
"""

from __future__ import annotations

import asyncio
import json
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from mcp_proxy_adapter.client.jsonrpc_client import JsonRpcClient
from mcp_proxy_adapter.client.jsonrpc_client.transport import JsonRpcTransport
from mcp_proxy_adapter.client.jsonrpc_client.ws_job_status import (
    WS_TERMINAL_EVENTS,
    BidirectionalWsChannel,
    open_bidirectional_ws_channel,
    wait_for_job_via_websocket,
)


def _make_mock_ws_message(msg_type: int, data: str | None = None) -> MagicMock:
    m = MagicMock()
    m.type = msg_type
    m.data = data
    return m


@pytest.mark.asyncio
async def test_wait_for_job_via_websocket_success() -> None:
    """wait_for_job_via_websocket returns payload when server sends job_completed."""
    transport = JsonRpcTransport(protocol="http", host="127.0.0.1", port=8080)
    job_id = "test-job-123"
    payload = {"event": "job_completed", "job_id": job_id, "result": {"done": True}}

    mock_ws = AsyncMock()
    mock_ws.send_str = AsyncMock()
    mock_ws.receive = AsyncMock(
        side_effect=[
            _make_mock_ws_message(1, json.dumps(payload)),  # aiohttp.WSMsgType.TEXT = 1
        ]
    )

    mock_ctx = AsyncMock()
    mock_ctx.__aenter__.return_value = mock_ws
    mock_ctx.__aexit__.return_value = None

    mock_session = MagicMock()
    mock_session.ws_connect = MagicMock(return_value=mock_ctx)
    mock_session_ctx = AsyncMock()
    mock_session_ctx.__aenter__.return_value = mock_session
    mock_session_ctx.__aexit__.return_value = None

    mock_session_class = MagicMock()
    mock_session_class.return_value = mock_session_ctx

    with patch("aiohttp.ClientSession", mock_session_class):
        result = await wait_for_job_via_websocket(transport, job_id, timeout=5.0)

    assert result == payload
    mock_ws.send_str.assert_called_once()
    call_arg = mock_ws.send_str.call_args[0][0]
    assert json.loads(call_arg) == {"action": "subscribe", "job_id": job_id}


@pytest.mark.asyncio
async def test_wait_for_job_via_websocket_job_failed() -> None:
    """wait_for_job_via_websocket returns payload when server sends job_failed."""
    transport = JsonRpcTransport(protocol="http", host="127.0.0.1", port=8080)
    job_id = "fail-job"
    payload = {"event": "job_failed", "job_id": job_id, "error": "Something broke"}

    mock_ws = AsyncMock()
    mock_ws.send_str = AsyncMock()
    mock_ws.receive = AsyncMock(
        return_value=_make_mock_ws_message(1, json.dumps(payload))
    )

    mock_ctx = AsyncMock()
    mock_ctx.__aenter__.return_value = mock_ws
    mock_ctx.__aexit__.return_value = None
    mock_session = MagicMock()
    mock_session.ws_connect = MagicMock(return_value=mock_ctx)
    mock_session_ctx = AsyncMock()
    mock_session_ctx.__aenter__.return_value = mock_session
    mock_session_ctx.__aexit__.return_value = None
    with patch("aiohttp.ClientSession", return_value=mock_session_ctx):
        result = await wait_for_job_via_websocket(transport, job_id, timeout=5.0)
    assert result["event"] == "job_failed"
    assert result["error"] == "Something broke"


@pytest.mark.asyncio
async def test_wait_for_job_via_websocket_timeout() -> None:
    """wait_for_job_via_websocket raises TimeoutError when no terminal event."""
    transport = JsonRpcTransport(protocol="http", host="127.0.0.1", port=8080)
    job_id = "slow-job"

    mock_ws = AsyncMock()
    mock_ws.send_str = AsyncMock()
    mock_ws.receive = AsyncMock(
        side_effect=asyncio.TimeoutError()
    )

    mock_ctx = AsyncMock()
    mock_ctx.__aenter__.return_value = mock_ws
    mock_ctx.__aexit__.return_value = None
    mock_session = MagicMock()
    mock_session.ws_connect = MagicMock(return_value=mock_ctx)
    mock_session_ctx = AsyncMock()
    mock_session_ctx.__aenter__.return_value = mock_session
    mock_session_ctx.__aexit__.return_value = None
    with patch("aiohttp.ClientSession", return_value=mock_session_ctx):
        with pytest.raises(asyncio.TimeoutError) as exc_info:
            await wait_for_job_via_websocket(transport, job_id, timeout=0.1)
    assert job_id in str(exc_info.value)


@pytest.mark.asyncio
async def test_wait_for_job_via_websocket_ws_close_raises() -> None:
    """wait_for_job_via_websocket raises RuntimeError when WebSocket closes."""
    import aiohttp

    transport = JsonRpcTransport(protocol="http", host="127.0.0.1", port=8080)
    job_id = "closed-job"

    mock_ws = AsyncMock()
    mock_ws.send_str = AsyncMock()
    mock_ws.receive = AsyncMock(
        return_value=_make_mock_ws_message(aiohttp.WSMsgType.CLOSE)
    )

    mock_ctx = AsyncMock()
    mock_ctx.__aenter__.return_value = mock_ws
    mock_ctx.__aexit__.return_value = None
    mock_session = MagicMock()
    mock_session.ws_connect = MagicMock(return_value=mock_ctx)
    mock_session_ctx = AsyncMock()
    mock_session_ctx.__aenter__.return_value = mock_session
    mock_session_ctx.__aexit__.return_value = None
    with patch("aiohttp.ClientSession", return_value=mock_session_ctx):
        with pytest.raises(RuntimeError) as exc_info:
            await wait_for_job_via_websocket(transport, job_id, timeout=5.0)
    assert "closed" in str(exc_info.value).lower() or "error" in str(exc_info.value).lower()


@pytest.mark.asyncio
async def test_wait_for_job_via_websocket_connection_error() -> None:
    """wait_for_job_via_websocket raises RuntimeError on ClientError."""
    transport = JsonRpcTransport(protocol="http", host="127.0.0.1", port=8080)
    job_id = "conn-fail"

    mock_session_ctx = AsyncMock()
    mock_session_ctx.__aenter__.side_effect = ConnectionError("refused")
    mock_session_ctx.__aexit__.return_value = None

    with patch("aiohttp.ClientSession", return_value=mock_session_ctx):
        with pytest.raises((RuntimeError, ConnectionError)):
            await wait_for_job_via_websocket(transport, job_id, timeout=5.0)


def test_ws_terminal_events() -> None:
    """WS_TERMINAL_EVENTS contains expected event names."""
    assert "job_completed" in WS_TERMINAL_EVENTS
    assert "job_failed" in WS_TERMINAL_EVENTS
    assert "job_stopped" in WS_TERMINAL_EVENTS


@pytest.mark.asyncio
async def test_bidirectional_channel_send_json_and_receive() -> None:
    """BidirectionalWsChannel send_json and receive_iter with mocked ws."""
    import aiohttp

    transport = JsonRpcTransport(protocol="http", host="127.0.0.1", port=8080)
    channel = BidirectionalWsChannel(transport, receive_timeout=5.0, heartbeat=10.0)

    mock_ws = AsyncMock()
    mock_ws.send_str = AsyncMock()
    received = [
        _make_mock_ws_message(aiohttp.WSMsgType.TEXT, json.dumps({"event": "progress", "p": 50})),
        _make_mock_ws_message(aiohttp.WSMsgType.TEXT, json.dumps({"event": "job_completed", "job_id": "j1"})),
        _make_mock_ws_message(aiohttp.WSMsgType.CLOSE),  # so receive_iter loop exits
    ]
    mock_ws.receive = AsyncMock(side_effect=received)
    mock_ws.close = AsyncMock()

    async def _ws_connect_return_ws(*args: object, **kwargs: object) -> MagicMock:
        return mock_ws

    mock_session_ctx = AsyncMock()
    mock_session_ctx.__aenter__.return_value = mock_session_ctx
    mock_session_ctx.__aexit__.return_value = None
    mock_session_ctx.close = AsyncMock()
    # Code does: self._ws = await self._session.ws_connect(...)
    mock_session_ctx.ws_connect = MagicMock(return_value=_ws_connect_return_ws())

    with patch("aiohttp.ClientSession", return_value=mock_session_ctx):
        async with channel:
            await channel.send_json({"action": "subscribe", "job_id": "j1"})
            messages = []
            async for msg in channel.receive_iter():
                messages.append(msg)
                if msg.get("event") == "job_completed":
                    break

    assert len(messages) == 2
    assert messages[0]["event"] == "progress"
    assert messages[1]["event"] == "job_completed"
    mock_ws.send_str.assert_called()
    call_arg = mock_ws.send_str.call_args_list[0][0][0]
    assert json.loads(call_arg) == {"action": "subscribe", "job_id": "j1"}


@pytest.mark.asyncio
async def test_bidirectional_channel_send_json_without_connect_raises() -> None:
    """BidirectionalWsChannel.send_json raises when not connected."""
    transport = JsonRpcTransport(protocol="http", host="127.0.0.1", port=8080)
    channel = BidirectionalWsChannel(transport)
    with pytest.raises(RuntimeError) as exc_info:
        await channel.send_json({"action": "subscribe", "job_id": "x"})
    assert "not connected" in str(exc_info.value).lower()


@pytest.mark.asyncio
async def test_bidirectional_channel_receive_iter_without_connect_raises() -> None:
    """BidirectionalWsChannel.receive_iter raises when not connected."""
    transport = JsonRpcTransport(protocol="http", host="127.0.0.1", port=8080)
    channel = BidirectionalWsChannel(transport)
    with pytest.raises(RuntimeError) as exc_info:
        async for _ in channel.receive_iter():
            pass
    assert "not connected" in str(exc_info.value).lower()


def test_open_bidirectional_ws_channel_params() -> None:
    """open_bidirectional_ws_channel passes receive_timeout and heartbeat."""
    transport = JsonRpcTransport(protocol="http", host="127.0.0.1", port=8080)
    ch = open_bidirectional_ws_channel(transport, receive_timeout=120.0, heartbeat=15.0)
    assert ch._receive_timeout == 120.0
    assert ch._heartbeat == 15.0
