"""Tests for execute_async WebSocket timeout bug fix.

Verifies that ws_id is correctly extracted from server responses (including
legacy deliver_id), validation prevents blank ws_id, and correct WebSocket
function is used.

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

from __future__ import annotations

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

import pytest

from mcp_proxy_adapter.client.jsonrpc_client.client import JsonRpcClient
from mcp_proxy_adapter.core.delivery.exceptions import InvalidDeliveryContractError


class TestExecuteAsyncWebSocketTimeoutFix:
    """Tests for WebSocket timeout bug fix in execute_async."""

    @pytest.mark.asyncio
    async def test_execute_async_uses_ws_id_from_acceptance(self) -> None:
        """execute_async correctly extracts ws_id from server response with ws_id field."""
        client = JsonRpcClient(host="127.0.0.1", port=8080)
        callback = MagicMock()
        ws_id = "ws-server-issued-123"

        async def fake_rpc(method: str, params: dict) -> dict:
            assert method == "execute_async"
            return {
                "jsonrpc": "2.0",
                "result": {
                    "accepted": True,
                    "ws_id": ws_id,
                    "delivery": {"mode": "ws", "terminal_only": True},
                },
                "id": 1,
            }

        with patch.object(
            client, "jsonrpc_call", new_callable=AsyncMock, side_effect=fake_rpc
        ):
            with patch(
                "mcp_proxy_adapter.client.jsonrpc_client.command_api.wait_for_terminal_delivery_via_websocket",
                new_callable=AsyncMock,
            ) as mock_wait:

                async def mock_delivery(*args: object, **kwargs: object) -> dict:
                    await asyncio.sleep(0.1)
                    return {"event": "delivery_completed", "ws_id": ws_id, "result": {}}

                mock_wait.side_effect = mock_delivery

                result = await client.execute_async("echo", {}, on_result=callback)
                assert result["ws_id"] == ws_id
                mock_wait.assert_called_once()
                # Verify ws_id was passed as second argument
                call_args = mock_wait.call_args
                assert call_args[0][1] == ws_id  # self is first arg, ws_id is second
                await asyncio.sleep(0.2)
                callback.assert_called_once()
                assert callback.call_args[0][0]["event"] == "delivery_completed"

    @pytest.mark.asyncio
    async def test_execute_async_handles_legacy_deliver_id(self) -> None:
        """execute_async correctly handles legacy deliver_id field (maps to ws_id)."""
        client = JsonRpcClient(host="127.0.0.1", port=8080)
        callback = MagicMock()
        deliver_id = "legacy-xyz-789"

        async def fake_rpc(method: str, params: dict) -> dict:
            return {
                "jsonrpc": "2.0",
                "result": {
                    "accepted": True,
                    "deliver_id": deliver_id,  # Legacy field, no ws_id
                    "delivery": {"mode": "ws", "terminal_only": True},
                },
                "id": 1,
            }

        with patch.object(
            client, "jsonrpc_call", new_callable=AsyncMock, side_effect=fake_rpc
        ):
            with patch(
                "mcp_proxy_adapter.client.jsonrpc_client.command_api.wait_for_terminal_delivery_via_websocket",
                new_callable=AsyncMock,
            ) as mock_wait:

                async def mock_delivery(*args: object, **kwargs: object) -> dict:
                    await asyncio.sleep(0.1)
                    # ws_id should be mapped from deliver_id
                    return {
                        "event": "delivery_completed",
                        "ws_id": deliver_id,
                        "result": {},
                    }

                mock_wait.side_effect = mock_delivery

                await client.execute_async("echo", {}, on_result=callback)
                # execute_async returns original result, but internally deliver_id is mapped to ws_id
                # Verify that wait_for_terminal_delivery_via_websocket was called with mapped ws_id
                mock_wait.assert_called_once()
                # Verify mapped deliver_id was passed as ws_id (second argument after self)
                call_args = mock_wait.call_args
                assert call_args[0][1] == deliver_id
                await asyncio.sleep(0.2)
                callback.assert_called_once()

    @pytest.mark.asyncio
    async def test_execute_async_fails_when_ws_id_missing(self) -> None:
        """execute_async raises error when ws_id is missing for WS delivery mode."""
        client = JsonRpcClient(host="127.0.0.1", port=8080)
        callback = MagicMock()

        async def fake_rpc(method: str, params: dict) -> dict:
            return {
                "jsonrpc": "2.0",
                "result": {
                    "accepted": True,
                    # Missing both ws_id and deliver_id
                    "delivery": {"mode": "ws", "terminal_only": True},
                },
                "id": 1,
            }

        with patch.object(
            client, "jsonrpc_call", new_callable=AsyncMock, side_effect=fake_rpc
        ):
            with patch(
                "mcp_proxy_adapter.client.jsonrpc_client.command_api.wait_for_terminal_delivery_via_websocket",
                new_callable=AsyncMock,
            ) as mock_wait:
                # parse_acceptance_payload raises ValueError when both ws_id and deliver_id are missing
                with pytest.raises(
                    (ValueError, InvalidDeliveryContractError),
                    match="(accepted response requires job_id or ws_id|missing ws_id)",
                ):
                    await client.execute_async("echo", {}, on_result=callback)
                # Verify no WebSocket task was created
                mock_wait.assert_not_called()
                # Verify callback was never invoked
                callback.assert_not_called()

    @pytest.mark.asyncio
    async def test_execute_async_uses_terminal_delivery_function(self) -> None:
        """execute_async always uses wait_for_terminal_delivery_via_websocket (not wait_for_job_via_websocket)."""
        client = JsonRpcClient(host="127.0.0.1", port=8080)
        callback = MagicMock()
        ws_id = "ws-test-456"

        async def fake_rpc(method: str, params: dict) -> dict:
            return {
                "jsonrpc": "2.0",
                "result": {
                    "accepted": True,
                    "ws_id": ws_id,
                    "delivery": {"mode": "ws", "terminal_only": True},
                },
                "id": 1,
            }

        with patch.object(
            client, "jsonrpc_call", new_callable=AsyncMock, side_effect=fake_rpc
        ):
            with patch(
                "mcp_proxy_adapter.client.jsonrpc_client.command_api.wait_for_terminal_delivery_via_websocket",
                new_callable=AsyncMock,
            ) as mock_terminal:
                with patch(
                    "mcp_proxy_adapter.client.jsonrpc_client.command_api.wait_for_job_via_websocket",
                    new_callable=AsyncMock,
                ) as mock_job:

                    async def mock_delivery(*args: object, **kwargs: object) -> dict:
                        await asyncio.sleep(0.1)
                        return {
                            "event": "delivery_completed",
                            "ws_id": ws_id,
                            "result": {},
                        }

                    mock_terminal.side_effect = mock_delivery

                    await client.execute_async("echo", {}, on_result=callback)
                    # Verify correct function was called
                    mock_terminal.assert_called_once()
                    assert mock_terminal.call_args[0][1] == ws_id
                    # Verify forbidden function was NOT called
                    mock_job.assert_not_called()
