"""Tests for ws_delivery terminal listener (subscribe by ws_id, terminal delivery only).

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

from __future__ import annotations

import asyncio
import json
from typing import Any, List
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from mcp_proxy_adapter.client.jsonrpc_client.transport import JsonRpcTransport
from mcp_proxy_adapter.client.jsonrpc_client.ws_job_status import (
    WS_DELIVERY_TERMINAL_EVENTS,
    wait_for_terminal_delivery_via_websocket,
)


def _make_transport() -> MagicMock:
    t = MagicMock(spec=JsonRpcTransport)
    t.ws_url = "ws://localhost:8080/ws"
    t.headers = {}
    t.get_ws_ssl_context.return_value = None
    return t


@pytest.mark.asyncio
async def test_subscribe_payload_uses_ws_delivery_and_id() -> None:
    sent: List[str] = []

    class FakeWs:
        async def send_str(self, data: str) -> None:
            sent.append(data)

        async def receive(self) -> Any:
            # Return terminal so loop exits
            msg = MagicMock()
            msg.type = 1  # TEXT
            msg.data = json.dumps(
                {
                    "event": "delivery_completed",
                    "ws_id": "ws-456",
                    "job_id": "job-123",
                    "result": {"success": True},
                    "error": None,
                }
            )
            return msg

        async def close(self) -> None:
            pass

    fake_ws = FakeWs()
    cm_ws = MagicMock(
        __aenter__=AsyncMock(return_value=fake_ws),
        __aexit__=AsyncMock(return_value=None),
    )
    session = MagicMock()
    session.ws_connect = MagicMock(return_value=cm_ws)
    session_cm = MagicMock(
        __aenter__=AsyncMock(return_value=session),
        __aexit__=AsyncMock(return_value=None),
    )

    with patch("aiohttp.ClientSession") as session_cls:
        session_cls.return_value = session_cm
        transport = _make_transport()
        result = await wait_for_terminal_delivery_via_websocket(
            transport, "ws-456", timeout=5.0
        )
        assert result.get("event") == "delivery_completed"
        assert result.get("ws_id") == "ws-456"
    assert len(sent) == 1
    payload = json.loads(sent[0])
    assert payload.get("action") == "subscribe"
    assert payload.get("subscription_type") == "ws_delivery"
    assert payload.get("ws_id") == "ws-456"


@pytest.mark.asyncio
async def test_returns_on_delivery_completed() -> None:
    call_count = 0

    async def fake_receive() -> Any:
        nonlocal call_count
        call_count += 1
        msg = MagicMock()
        msg.type = 1
        msg.data = json.dumps(
            {
                "event": "delivery_completed",
                "ws_id": "ws-1",
                "result": {"success": True},
                "error": None,
            }
        )
        return msg

    fake_ws = MagicMock()
    fake_ws.send_str = AsyncMock()
    fake_ws.receive = fake_receive
    fake_ws.close = AsyncMock()

    session = MagicMock()
    session.ws_connect = MagicMock(
        return_value=MagicMock(
            __aenter__=AsyncMock(return_value=fake_ws),
            __aexit__=AsyncMock(return_value=None),
        )
    )
    session_cm = MagicMock(
        __aenter__=AsyncMock(return_value=session),
        __aexit__=AsyncMock(return_value=None),
    )
    with patch("aiohttp.ClientSession", return_value=session_cm):
        transport = _make_transport()
        result = await wait_for_terminal_delivery_via_websocket(
            transport, "ws-1", timeout=10.0
        )
    assert result["event"] == "delivery_completed"
    assert result["ws_id"] == "ws-1"


@pytest.mark.asyncio
async def test_timeout() -> None:
    async def never_terminal() -> Any:
        await asyncio.sleep(10)
        msg = MagicMock()
        msg.type = 1
        msg.data = json.dumps({"event": "delivery_pending", "ws_id": "ws-1"})
        return msg

    fake_ws = MagicMock()
    fake_ws.send_str = AsyncMock()
    fake_ws.receive = never_terminal

    session = MagicMock()
    session.ws_connect = MagicMock(
        return_value=MagicMock(
            __aenter__=AsyncMock(return_value=fake_ws),
            __aexit__=AsyncMock(return_value=None),
        )
    )
    session_cm = MagicMock(
        __aenter__=AsyncMock(return_value=session),
        __aexit__=AsyncMock(return_value=None),
    )
    with patch("aiohttp.ClientSession", return_value=session_cm):
        transport = _make_transport()
        with pytest.raises(asyncio.TimeoutError) as exc_info:
            await wait_for_terminal_delivery_via_websocket(
                transport, "ws-1", timeout=0.05
            )
        assert (
            "ws_id" in str(exc_info.value) or "terminal" in str(exc_info.value).lower()
        )


@pytest.mark.asyncio
async def test_ignores_non_terminal_events() -> None:
    events = [
        {"event": "delivery_pending", "ws_id": "ws-1"},
        {"event": "delivery_progress", "ws_id": "ws-1", "progress": 50},
        {"event": "delivery_completed", "ws_id": "ws-1", "result": {}},
    ]
    idx = 0

    async def next_msg() -> Any:
        nonlocal idx
        if idx >= len(events):
            await asyncio.sleep(10)
            msg = MagicMock()
            msg.type = 1
            msg.data = "{}"
            return msg
        msg = MagicMock()
        msg.type = 1
        msg.data = json.dumps(events[idx])
        idx += 1
        return msg

    fake_ws = MagicMock()
    fake_ws.send_str = AsyncMock()
    fake_ws.receive = next_msg

    session = MagicMock()
    session.ws_connect = MagicMock(
        return_value=MagicMock(
            __aenter__=AsyncMock(return_value=fake_ws),
            __aexit__=AsyncMock(return_value=None),
        )
    )
    session_cm = MagicMock(
        __aenter__=AsyncMock(return_value=session),
        __aexit__=AsyncMock(return_value=None),
    )
    with patch("aiohttp.ClientSession", return_value=session_cm):
        transport = _make_transport()
        result = await wait_for_terminal_delivery_via_websocket(
            transport, "ws-1", timeout=5.0
        )
    assert result["event"] == "delivery_completed"


def test_ws_delivery_terminal_events_constant() -> None:
    assert "delivery_completed" in WS_DELIVERY_TERMINAL_EVENTS
    assert "delivery_failed" in WS_DELIVERY_TERMINAL_EVENTS
    assert "delivery_cancelled" in WS_DELIVERY_TERMINAL_EVENTS
    assert "job_completed" not in WS_DELIVERY_TERMINAL_EVENTS
