"""
Tests for concurrent MCP tool calls and MCPClientManager.

These tests verify that multiple concurrent tool calls work correctly
without race conditions or deadlocks.

Tests are organized into:
1. Mock-based tests for fast, reliable testing of concurrency logic
2. Subprocess tests (marked with `integration`) for real-world verification

To run the subprocess tests locally:
    pytest tests/mcp/test_concurrent_mcp_calls.py -m integration
"""

import asyncio
import sys
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

# Skip all tests if MCP is not available
try:
    import mcp  # noqa: F401

    MCP_AVAILABLE = True
except ImportError:
    MCP_AVAILABLE = False

pytestmark = pytest.mark.skipif(
    not MCP_AVAILABLE, reason="MCP package not installed"
)

# Path to the echo server for testing
ECHO_SERVER_PATH = Path(__file__).parent / "echo_server.py"


class TestConcurrentToolCallsWithMocks:
    """Tests for concurrent tool calls using mocks for reliable testing."""

    @pytest.fixture
    def mock_tool_result(self):
        """Create a mock tool result factory."""

        def _make_result(text_content, is_error=False):
            mock_item = MagicMock()
            mock_item.text = text_content
            mock_result = MagicMock()
            mock_result.isError = is_error
            mock_result.content = [mock_item]
            return mock_result

        return _make_result

    @pytest.mark.asyncio
    async def test_parallel_tool_calls_with_mock(self, mock_tool_result):
        """Test multiple tool calls running in parallel."""
        from local_deep_research.mcp.client import MCPClient

        with patch("local_deep_research.mcp.client.MCP_AVAILABLE", True):
            config = {"name": "mock-server", "command": "python"}
            client = MCPClient(config)

            client._connected = True
            client._session = MagicMock()

            call_order = []

            async def mock_call(name, args):
                call_order.append(f"start_{args.get('id', 0)}")
                await asyncio.sleep(0.01)  # Simulate async work
                call_order.append(f"end_{args.get('id', 0)}")
                return mock_tool_result(f"Result for {args.get('id', 0)}")

            client._session.call_tool = mock_call

            # Launch concurrent calls
            tasks = [client.call_tool("echo", {"id": i}) for i in range(5)]

            results = await asyncio.gather(*tasks)

            assert len(results) == 5
            for result in results:
                assert result["status"] == "success"

    @pytest.mark.asyncio
    async def test_concurrent_calls_no_data_corruption_with_mock(
        self, mock_tool_result
    ):
        """Test that concurrent calls don't corrupt each other's data."""
        from local_deep_research.mcp.client import MCPClient

        with patch("local_deep_research.mcp.client.MCP_AVAILABLE", True):
            config = {"name": "mock-server", "command": "python"}
            client = MCPClient(config)

            client._connected = True
            client._session = MagicMock()

            # Track what was received
            received_messages = []

            async def mock_call(name, args):
                msg = args.get("message", "")
                received_messages.append(msg)
                await asyncio.sleep(0.001)  # Tiny delay
                return mock_tool_result(f'{{"echoed": "{msg}"}}')

            client._session.call_tool = mock_call

            # Create unique messages
            messages = [f"UNIQUE_{i}_MSG" for i in range(20)]
            tasks = [
                client.call_tool("echo", {"message": msg}) for msg in messages
            ]

            results = await asyncio.gather(*tasks)

            # All messages should have been received
            assert len(received_messages) == 20
            assert set(received_messages) == set(messages)

            # All results should succeed
            for result in results:
                assert result["status"] == "success"

    @pytest.mark.asyncio
    async def test_mixed_parallel_tools_with_mock(self, mock_tool_result):
        """Test different tools called in parallel."""
        from local_deep_research.mcp.client import MCPClient

        with patch("local_deep_research.mcp.client.MCP_AVAILABLE", True):
            config = {"name": "mock-server", "command": "python"}
            client = MCPClient(config)

            client._connected = True
            client._session = MagicMock()

            async def mock_call(name, args):
                await asyncio.sleep(0.001)
                if name == "echo":
                    return mock_tool_result(
                        f'{{"message": "{args.get("message", "")}"}}'
                    )
                if name == "add":
                    result = args.get("a", 0) + args.get("b", 0)
                    return mock_tool_result(f'{{"result": {result}}}')
                return mock_tool_result('{"status": "unknown"}')

            client._session.call_tool = mock_call

            tasks = [
                client.call_tool("echo", {"message": "hello"}),
                client.call_tool("add", {"a": 1, "b": 2}),
                client.call_tool("echo", {"message": "world"}),
                client.call_tool("add", {"a": 10, "b": 20}),
            ]

            results = await asyncio.gather(*tasks)

            assert len(results) == 4
            assert all(r["status"] == "success" for r in results)


class TestMCPClientManagerWithMocks:
    """Tests for MCPClientManager using mocks."""

    @pytest.mark.asyncio
    async def test_manager_call_tool_on_unknown_server(self):
        """Test calling tool on non-existent server raises error."""
        from local_deep_research.mcp.client import (
            MCPClientManager,
            MCPClientError,
        )

        manager = MCPClientManager([])

        async with manager.connect_all():
            with pytest.raises(MCPClientError) as exc_info:
                await manager.call_tool("unknown-server", "echo", {})

            assert "not connected" in str(exc_info.value)

    @pytest.mark.asyncio
    async def test_manager_empty_configs(self):
        """Test manager with empty config list."""
        from local_deep_research.mcp.client import MCPClientManager

        manager = MCPClientManager([])

        async with manager.connect_all():
            servers = manager.get_connected_servers()
            assert servers == []

    @pytest.mark.asyncio
    async def test_manager_cleanup_tracking(self):
        """Test that manager properly tracks cleanup."""
        from local_deep_research.mcp.client import MCPClientManager

        # Use empty configs to avoid actual connections
        manager = MCPClientManager([])

        assert len(manager._clients) == 0

        async with manager.connect_all():
            assert len(manager._clients) == 0  # No servers configured

        assert len(manager._clients) == 0


# =============================================================================
# Subprocess-based integration tests - marked for opt-in execution
# =============================================================================


@pytest.mark.integration
class TestConcurrentToolCallsSubprocess:
    """Tests for concurrent tool calls using real subprocess."""

    @pytest.fixture
    def echo_server_config(self):
        """Configuration for the echo test server."""
        return {
            "name": "echo-concurrent-test",
            "command": sys.executable,
            "args": [str(ECHO_SERVER_PATH)],
        }

    @pytest.mark.asyncio
    async def test_parallel_tool_calls(self, echo_server_config):
        """Test multiple tool calls running in parallel."""
        from local_deep_research.mcp.client import MCPClient

        client = MCPClient(echo_server_config, timeout=60.0)

        async with client.connect() as connected:
            # Launch multiple tool calls concurrently
            tasks = [
                connected.call_tool("echo", {"message": f"Message {i}"})
                for i in range(5)
            ]

            results = await asyncio.gather(*tasks)

            # All should succeed
            assert len(results) == 5
            for i, result in enumerate(results):
                assert result["status"] == "success"
                assert f"Message {i}" in result["content"]

    @pytest.mark.asyncio
    async def test_mixed_parallel_tool_calls(self, echo_server_config):
        """Test different tools called in parallel."""
        from local_deep_research.mcp.client import MCPClient

        client = MCPClient(echo_server_config, timeout=60.0)

        async with client.connect() as connected:
            # Launch different tools concurrently
            tasks = [
                connected.call_tool("echo", {"message": "test"}),
                connected.call_tool("add_numbers", {"a": 1, "b": 2}),
                connected.call_tool("get_info", {}),
                connected.call_tool("add_numbers", {"a": 10, "b": 20}),
            ]

            results = await asyncio.gather(*tasks)

            # All should succeed
            assert len(results) == 4
            for result in results:
                assert result["status"] == "success"

            # Verify specific results
            assert "test" in results[0]["content"]
            assert "3" in results[1]["content"]
            assert "echo-test-server" in results[2]["content"]
            assert "30" in results[3]["content"]

    @pytest.mark.asyncio
    async def test_concurrent_calls_no_data_corruption(
        self, echo_server_config
    ):
        """Test that concurrent calls don't corrupt each other's data."""
        from local_deep_research.mcp.client import MCPClient

        client = MCPClient(echo_server_config, timeout=60.0)

        async with client.connect() as connected:
            # Create unique messages to verify no cross-contamination
            messages = [f"UNIQUE_{i}_MESSAGE_{i * 100}" for i in range(10)]

            tasks = [
                connected.call_tool("echo", {"message": msg})
                for msg in messages
            ]

            results = await asyncio.gather(*tasks)

            # Each result should contain its corresponding message
            for i, (msg, result) in enumerate(zip(messages, results)):
                assert result["status"] == "success"
                assert msg in result["content"], f"Message {i} was corrupted"


@pytest.mark.integration
class TestMCPClientManagerSubprocess:
    """Tests for MCPClientManager with real subprocess servers."""

    @pytest.mark.asyncio
    async def test_manager_connect_all_single_server(self):
        """Test MCPClientManager connecting to a single server."""
        from local_deep_research.mcp.client import MCPClientManager

        configs = [
            {
                "name": "echo-1",
                "command": sys.executable,
                "args": [str(ECHO_SERVER_PATH)],
            }
        ]

        manager = MCPClientManager(configs)

        async with manager.connect_all():
            servers = manager.get_connected_servers()
            assert "echo-1" in servers

            # Test calling a tool
            result = await manager.call_tool(
                "echo-1", "echo", {"message": "Manager test"}
            )
            assert result["status"] == "success"
            assert "Manager test" in result["content"]

    @pytest.mark.asyncio
    async def test_manager_list_all_tools(self):
        """Test listing tools from all connected servers."""
        from local_deep_research.mcp.client import MCPClientManager

        configs = [
            {
                "name": "echo-tools",
                "command": sys.executable,
                "args": [str(ECHO_SERVER_PATH)],
            }
        ]

        manager = MCPClientManager(configs)

        async with manager.connect_all():
            all_tools = await manager.list_all_tools()

            assert "echo-tools" in all_tools
            tool_names = [t["name"] for t in all_tools["echo-tools"]]
            assert "echo" in tool_names
            assert "add_numbers" in tool_names

    @pytest.mark.asyncio
    async def test_manager_cleanup_on_exit(self):
        """Test that manager properly cleans up connections."""
        from local_deep_research.mcp.client import MCPClientManager

        configs = [
            {
                "name": "cleanup-test",
                "command": sys.executable,
                "args": [str(ECHO_SERVER_PATH)],
            }
        ]

        manager = MCPClientManager(configs)

        async with manager.connect_all():
            assert len(manager.get_connected_servers()) == 1

        # After exit, should have no connected servers
        assert len(manager._clients) == 0

    @pytest.mark.asyncio
    async def test_manager_handles_connection_failure_gracefully(self):
        """Test that manager continues when one server fails to connect."""
        from local_deep_research.mcp.client import MCPClientManager

        configs = [
            {
                "name": "good-server",
                "command": sys.executable,
                "args": [str(ECHO_SERVER_PATH)],
            },
            {
                "name": "bad-server",
                "command": "python",
                "args": ["/nonexistent/path.py"],
            },
        ]

        manager = MCPClientManager(configs)

        async with manager.connect_all():
            servers = manager.get_connected_servers()
            # The good server should be connected
            assert "good-server" in servers
            # The bad server should NOT be connected (failed gracefully)
            assert "bad-server" not in servers


class TestConcurrencyMocking:
    """Tests for concurrent behavior using mocks (for faster/deterministic testing)."""

    @pytest.mark.asyncio
    async def test_concurrent_calls_with_mocked_session(self):
        """Test concurrent behavior with mocked MCP session."""
        from local_deep_research.mcp.client import MCPClient

        # Create a mock that simulates some async delay
        async def mock_call_tool(name, args):
            await asyncio.sleep(0.01)  # Small delay to test concurrency
            mock_result = MagicMock()
            mock_result.isError = False
            mock_item = MagicMock()
            mock_item.text = f"Result for {name} with {args}"
            mock_result.content = [mock_item]
            return mock_result

        with patch("local_deep_research.mcp.client.MCP_AVAILABLE", True):
            config = {"name": "mock-server", "command": "python"}
            client = MCPClient(config)

            # Manually set up the client as if connected
            client._connected = True
            client._session = MagicMock()
            client._session.call_tool = mock_call_tool

            # Run concurrent calls
            tasks = [client.call_tool("tool", {"id": i}) for i in range(10)]

            results = await asyncio.gather(*tasks)

            assert len(results) == 10
            for result in results:
                assert result["status"] == "success"

    @pytest.mark.asyncio
    async def test_no_race_condition_on_connection_state(self):
        """Test that connection state is properly synchronized."""
        from local_deep_research.mcp.client import MCPClient, MCPClientError

        with patch("local_deep_research.mcp.client.MCP_AVAILABLE", True):
            config = {"name": "race-test", "command": "python"}
            client = MCPClient(config)

            # Not connected - should raise
            async def attempt_call():
                return await client.call_tool("test", {})

            tasks = [attempt_call() for _ in range(5)]

            # All should raise MCPClientError since not connected
            results = await asyncio.gather(*tasks, return_exceptions=True)

            for result in results:
                assert isinstance(result, MCPClientError)
                assert "Not connected" in str(result)


class TestAsyncUtilityConcurrency:
    """Tests for the run_async utility with concurrent scenarios."""

    def test_run_async_multiple_sequential_calls(self):
        """Test run_async handles multiple sequential calls."""
        from local_deep_research.mcp.client import run_async

        async def coro(n):
            return n * 2

        results = [run_async(coro(i)) for i in range(5)]

        assert results == [0, 2, 4, 6, 8]

    def test_run_async_with_exception_in_one(self):
        """Test run_async propagates exceptions correctly."""
        from local_deep_research.mcp.client import run_async

        async def success_coro():
            return "ok"

        async def fail_coro():
            raise ValueError("failed")

        # First call succeeds
        assert run_async(success_coro()) == "ok"

        # Second call fails
        with pytest.raises(ValueError) as exc_info:
            run_async(fail_coro())
        assert "failed" in str(exc_info.value)

        # Third call succeeds (exception doesn't affect future calls)
        assert run_async(success_coro()) == "ok"
