"""Tests for execute_command_unified method.

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

import asyncio
import pytest
from unittest.mock import AsyncMock, patch
from typing import Any, Dict, List

from mcp_proxy_adapter.client.jsonrpc_client import JsonRpcClient, QueueJobStatus

# Patch target for wait_for_job_via_websocket (used when auto_poll=True and job_id)
WS_PATCH = (
    "mcp_proxy_adapter.client.jsonrpc_client.command_api.wait_for_job_via_websocket"
)


class TestExecuteCommandUnified:
    """Test suite for execute_command_unified method."""

    @pytest.fixture
    def client(self) -> JsonRpcClient:
        """Create test client instance."""
        return JsonRpcClient(protocol="http", host="127.0.0.1", port=8080)

    @pytest.mark.asyncio
    async def test_immediate_execution_no_job_id(self, client: JsonRpcClient) -> None:
        """Test immediate execution when no job_id is returned."""
        # Mock execute_command to return immediate result
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "data": {"result": "test"}}
        )

        result = await client.execute_command_unified("echo", {"message": "test"})

        assert result["mode"] == "immediate"
        assert result["queued"] is False
        assert result["command"] == "echo"
        assert "result" in result

    @pytest.mark.asyncio
    async def test_immediate_execution_expect_queue_false(
        self, client: JsonRpcClient
    ) -> None:
        """Test immediate execution when expect_queue=False even with job_id."""
        # Mock execute_command to return job_id
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "job_id": "test-job-123"}
        )

        result = await client.execute_command_unified(
            "echo", {"message": "test"}, expect_queue=False
        )

        assert result["mode"] == "immediate"
        assert result["queued"] is False

    @pytest.mark.asyncio
    async def test_expect_queue_true_no_job_id_raises_error(
        self, client: JsonRpcClient
    ) -> None:
        """Test that expect_queue=True without job_id raises RuntimeError."""
        # Mock execute_command to return no job_id
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "data": {"result": "test"}}
        )

        with pytest.raises(RuntimeError, match="expected to run via queue"):
            await client.execute_command_unified(
                "long_task", {"seconds": 10}, expect_queue=True
            )

    @pytest.mark.asyncio
    async def test_queued_execution_auto_poll_false(
        self, client: JsonRpcClient
    ) -> None:
        """Test queued execution with auto_poll=False."""
        job_id = "test-job-123"
        # Mock execute_command to return job_id
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "job_id": job_id}
        )

        result = await client.execute_command_unified(
            "long_task", {"seconds": 10}, auto_poll=False
        )

        assert result["mode"] == "queued"
        assert result["queued"] is True
        assert result["job_id"] == job_id
        assert "status" in result

    @pytest.mark.asyncio
    async def test_queued_execution_auto_poll_true_completed(
        self, client: JsonRpcClient
    ) -> None:
        """Test queued execution with auto_poll=True and completed status (WS)."""
        job_id = "test-job-123"
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "job_id": job_id}
        )
        with patch(WS_PATCH, new_callable=AsyncMock) as mock_ws:
            mock_ws.return_value = {
                "event": "job_completed",
                "result": {"output": "done"},
            }
            result = await client.execute_command_unified(
                "long_task", {"seconds": 10}, auto_poll=True, poll_interval=0.1
            )
        assert result["mode"] == "queued"
        assert result["queued"] is True
        assert result["job_id"] == job_id
        assert result["status"] == "job_completed"
        assert result["result"] == {"output": "done"}

    @pytest.mark.asyncio
    async def test_queued_execution_auto_poll_true_failed(
        self, client: JsonRpcClient
    ) -> None:
        """Test queued execution with auto_poll=True and failed status (WS)."""
        job_id = "test-job-123"
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "job_id": job_id}
        )
        with patch(WS_PATCH, new_callable=AsyncMock) as mock_ws:
            mock_ws.return_value = {"event": "job_failed", "error": "Task failed"}
            with pytest.raises(RuntimeError, match="failed"):
                await client.execute_command_unified(
                    "long_task", {"seconds": 10}, auto_poll=True, poll_interval=0.1
                )

    @pytest.mark.asyncio
    async def test_queued_execution_timeout(self, client: JsonRpcClient) -> None:
        """Test queued execution with timeout (WS raises TimeoutError)."""
        job_id = "test-job-123"
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "job_id": job_id}
        )
        with patch(WS_PATCH, new_callable=AsyncMock) as mock_ws:
            mock_ws.side_effect = asyncio.TimeoutError("Job did not finish within 0.3s")
            with pytest.raises(
                (TimeoutError, asyncio.TimeoutError), match="did not finish within"
            ):
                await client.execute_command_unified(
                    "long_task",
                    {"seconds": 10},
                    auto_poll=True,
                    poll_interval=0.1,
                    timeout=0.3,
                )

    @pytest.mark.asyncio
    async def test_status_hook_invocation(self, client: JsonRpcClient) -> None:
        """With WS, status_hook is unused (WS delivers final event only)."""
        job_id = "test-job-123"
        status_hook_calls: List[Dict[str, Any]] = []

        async def status_hook(status: Dict[str, Any]) -> None:
            status_hook_calls.append(status)

        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "job_id": job_id}
        )
        with patch(WS_PATCH, new_callable=AsyncMock) as mock_ws:
            mock_ws.return_value = {
                "event": "job_completed",
                "result": {"output": "done"},
            }
            await client.execute_command_unified(
                "long_task",
                {"seconds": 10},
                auto_poll=True,
                poll_interval=0.1,
                status_hook=status_hook,
            )
        # WS path does not invoke status_hook (docstring says unused).
        assert len(status_hook_calls) == 0

    @pytest.mark.asyncio
    async def test_cancelled_status_handling(self, client: JsonRpcClient) -> None:
        """Test handling of cancelled/stopped status from WS event."""
        job_id = "test-job-123"
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "job_id": job_id}
        )
        with patch(WS_PATCH, new_callable=AsyncMock) as mock_ws:
            mock_ws.return_value = {"event": "job_stopped", "result": None}
            result = await client.execute_command_unified(
                "long_task", {"seconds": 10}, auto_poll=True, poll_interval=0.1
            )
        assert result["mode"] == "queued"
        assert result["status"] == "job_stopped"
        assert result["result"] is None

    @pytest.mark.asyncio
    async def test_unknown_status_handling(self, client: JsonRpcClient) -> None:
        """Test handling of unknown/non-failure event type from WS."""
        job_id = "test-job-123"
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "job_id": job_id}
        )
        with patch(WS_PATCH, new_callable=AsyncMock) as mock_ws:
            mock_ws.return_value = {
                "event": "unknown_status",
                "result": None,
            }
            result = await client.execute_command_unified(
                "long_task", {"seconds": 10}, auto_poll=True, poll_interval=0.1
            )
        assert result["mode"] == "queued"
        assert result["status"] == "unknown_status"
        assert result["result"] is None

    @pytest.mark.asyncio
    async def test_queue_get_job_status_error_handling(
        self, client: JsonRpcClient
    ) -> None:
        """Test error when wait_for_job_via_websocket fails (e.g. connection)."""
        job_id = "test-job-123"
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "job_id": job_id}
        )
        with patch(WS_PATCH, new_callable=AsyncMock) as mock_ws:
            mock_ws.side_effect = RuntimeError("Failed to get status")
            with pytest.raises(RuntimeError, match="Failed to get status"):
                await client.execute_command_unified(
                    "long_task", {"seconds": 10}, auto_poll=True, poll_interval=0.1
                )

    @pytest.mark.asyncio
    async def test_nested_result_extraction(self, client: JsonRpcClient) -> None:
        """Test extraction of result from nested structures (WS event)."""
        job_id = "test-job-123"
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "job_id": job_id}
        )
        with patch(WS_PATCH, new_callable=AsyncMock) as mock_ws:
            mock_ws.return_value = {
                "event": "job_completed",
                "result": {"data": {"output": "nested_result"}},
            }
            result = await client.execute_command_unified(
                "long_task", {"seconds": 10}, auto_poll=True, poll_interval=0.1
            )
        assert result["result"] == {"output": "nested_result"}

    @pytest.mark.asyncio
    async def test_use_cmd_endpoint(self, client: JsonRpcClient) -> None:
        """Test execution using /cmd endpoint."""
        # Mock cmd_call instead of jsonrpc_call
        client.cmd_call = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "data": {"result": "test"}}
        )

        result = await client.execute_command_unified(
            "echo", {"message": "test"}, use_cmd_endpoint=True
        )

        assert result["mode"] == "immediate"
        client.cmd_call.assert_called_once()

    @pytest.mark.asyncio
    async def test_job_id_extraction_from_data_section(
        self, client: JsonRpcClient
    ) -> None:
        """Test extraction of job_id from data section; WS returns result."""
        job_id = "test-job-123"
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "data": {"jobId": job_id}}
        )
        with patch(WS_PATCH, new_callable=AsyncMock) as mock_ws:
            mock_ws.return_value = {
                "event": "job_completed",
                "result": {"output": "done"},
            }
            result = await client.execute_command_unified(
                "long_task", {"seconds": 10}, auto_poll=True, poll_interval=0.1
            )
        assert result["mode"] == "queued"
        assert result["job_id"] == job_id

    @pytest.mark.asyncio
    async def test_all_pending_states_polling(self, client: JsonRpcClient) -> None:
        """With WS, each run gets one completion event; no polling."""
        job_id = "test-job-123"
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "job_id": job_id}
        )
        ws_event = {"event": "job_completed", "result": {"output": "done"}}
        with patch(WS_PATCH, new_callable=AsyncMock) as mock_ws:
            mock_ws.return_value = ws_event
            for _ in QueueJobStatus.get_pending_states():
                result = await client.execute_command_unified(
                    "long_task", {"seconds": 10}, auto_poll=True, poll_interval=0.1
                )
                assert result["status"] == "job_completed"

    @pytest.mark.asyncio
    async def test_all_failure_states_raise_error(self, client: JsonRpcClient) -> None:
        """Test that job_failed WS event raises RuntimeError."""
        job_id = "test-job-123"
        client.execute_command = AsyncMock(  # type: ignore[method-assign]
            return_value={"success": True, "job_id": job_id}
        )
        with patch(WS_PATCH, new_callable=AsyncMock) as mock_ws:
            mock_ws.return_value = {"event": "job_failed", "error": "Failed"}
            with pytest.raises(RuntimeError, match="failed"):
                await client.execute_command_unified(
                    "long_task", {"seconds": 10}, auto_poll=True, poll_interval=0.1
                )
