import contextlib
from unittest import mock

import httpx
import mcp_types as types
import pytest

import mcp
from mcp.client.session_group import (
    ClientSessionGroup,
    ClientSessionParameters,
    SseServerParameters,
    StreamableHttpParameters,
)
from mcp.client.stdio import StdioServerParameters
from mcp.shared.exceptions import MCPError


@pytest.fixture
def mock_exit_stack():
    """Fixture for a mocked AsyncExitStack."""
    # Use unittest.mock.Mock directly if needed, or just a plain object
    # if only attribute access/existence is needed.
    # For AsyncExitStack, Mock or MagicMock is usually fine.
    return mock.MagicMock(spec=contextlib.AsyncExitStack)


def test_client_session_group_init():
    mcp_session_group = ClientSessionGroup()
    assert not mcp_session_group._tools
    assert not mcp_session_group._resources
    assert not mcp_session_group._prompts
    assert not mcp_session_group._tool_to_session


def test_client_session_group_component_properties():
    # --- Mock Dependencies ---
    mock_prompt = mock.Mock()
    mock_resource = mock.Mock()
    mock_tool = mock.Mock()

    # --- Prepare Session Group ---
    mcp_session_group = ClientSessionGroup()
    mcp_session_group._prompts = {"my_prompt": mock_prompt}
    mcp_session_group._resources = {"my_resource": mock_resource}
    mcp_session_group._tools = {"my_tool": mock_tool}

    # --- Assertions ---
    assert mcp_session_group.prompts == {"my_prompt": mock_prompt}
    assert mcp_session_group.resources == {"my_resource": mock_resource}
    assert mcp_session_group.tools == {"my_tool": mock_tool}


@pytest.mark.anyio
async def test_client_session_group_call_tool():
    # --- Mock Dependencies ---
    mock_session = mock.AsyncMock()

    # --- Prepare Session Group ---
    def hook(name: str, server_info: types.Implementation) -> str:  # pragma: no cover
        return f"{(server_info.name)}-{name}"

    mcp_session_group = ClientSessionGroup(component_name_hook=hook)
    mcp_session_group._tools = {"server1-my_tool": types.Tool(name="my_tool", input_schema={})}
    mcp_session_group._tool_to_session = {"server1-my_tool": mock_session}
    text_content = types.TextContent(type="text", text="OK")
    mock_session.call_tool.return_value = types.CallToolResult(content=[text_content])

    # --- Test Execution ---
    result = await mcp_session_group.call_tool(
        name="server1-my_tool",
        arguments={
            "name": "value1",
            "args": {},
        },
    )

    # --- Assertions ---
    assert result.content == [text_content]
    mock_session.call_tool.assert_called_once_with(
        "my_tool",
        arguments={"name": "value1", "args": {}},
        read_timeout_seconds=None,
        progress_callback=None,
        input_responses=None,
        request_state=None,
        meta=None,
        allow_input_required=False,
    )


@pytest.mark.anyio
async def test_client_session_group_call_tool_forwards_allow_input_required():
    mock_session = mock.AsyncMock()
    mcp_session_group = ClientSessionGroup()
    mcp_session_group._tools = {"my_tool": types.Tool(name="my_tool", input_schema={})}
    mcp_session_group._tool_to_session = {"my_tool": mock_session}
    mock_session.call_tool.return_value = types.InputRequiredResult(request_state="s")

    result = await mcp_session_group.call_tool(name="my_tool", arguments={}, allow_input_required=True)
    assert isinstance(result, types.InputRequiredResult)
    assert result.request_state == "s"
    assert mock_session.call_tool.call_args.kwargs["allow_input_required"] is True


@pytest.mark.anyio
async def test_client_session_group_connect_to_server(mock_exit_stack: contextlib.AsyncExitStack):
    """Test connecting to a server and aggregating components."""
    # --- Mock Dependencies ---
    mock_server_info = mock.Mock(spec=types.Implementation)
    mock_server_info.name = "TestServer1"
    mock_session = mock.AsyncMock(spec=mcp.ClientSession)
    mock_tool1 = mock.Mock(spec=types.Tool)
    mock_tool1.name = "tool_a"
    mock_resource1 = mock.Mock(spec=types.Resource)
    mock_resource1.name = "resource_b"
    mock_prompt1 = mock.Mock(spec=types.Prompt)
    mock_prompt1.name = "prompt_c"
    mock_session.list_tools.return_value = mock.AsyncMock(tools=[mock_tool1])
    mock_session.list_resources.return_value = mock.AsyncMock(resources=[mock_resource1])
    mock_session.list_prompts.return_value = mock.AsyncMock(prompts=[mock_prompt1])

    # --- Test Execution ---
    group = ClientSessionGroup(exit_stack=mock_exit_stack)
    with mock.patch.object(group, "_establish_session", return_value=(mock_server_info, mock_session)):
        await group.connect_to_server(StdioServerParameters(command="test"))

    # --- Assertions ---
    assert mock_session in group._sessions
    assert len(group.tools) == 1
    assert "tool_a" in group.tools
    assert group.tools["tool_a"] == mock_tool1
    assert group._tool_to_session["tool_a"] == mock_session
    assert len(group.resources) == 1
    assert "resource_b" in group.resources
    assert group.resources["resource_b"] == mock_resource1
    assert len(group.prompts) == 1
    assert "prompt_c" in group.prompts
    assert group.prompts["prompt_c"] == mock_prompt1
    mock_session.list_tools.assert_awaited_once()
    mock_session.list_resources.assert_awaited_once()
    mock_session.list_prompts.assert_awaited_once()


@pytest.mark.anyio
async def test_client_session_group_connect_to_server_with_name_hook(mock_exit_stack: contextlib.AsyncExitStack):
    """Test connecting with a component name hook."""
    # --- Mock Dependencies ---
    mock_server_info = mock.Mock(spec=types.Implementation)
    mock_server_info.name = "HookServer"
    mock_session = mock.AsyncMock(spec=mcp.ClientSession)
    mock_tool = mock.Mock(spec=types.Tool)
    mock_tool.name = "base_tool"
    mock_session.list_tools.return_value = mock.AsyncMock(tools=[mock_tool])
    mock_session.list_resources.return_value = mock.AsyncMock(resources=[])
    mock_session.list_prompts.return_value = mock.AsyncMock(prompts=[])

    # --- Test Setup ---
    def name_hook(name: str, server_info: types.Implementation) -> str:
        return f"{server_info.name}.{name}"

    # --- Test Execution ---
    group = ClientSessionGroup(exit_stack=mock_exit_stack, component_name_hook=name_hook)
    with mock.patch.object(group, "_establish_session", return_value=(mock_server_info, mock_session)):
        await group.connect_to_server(StdioServerParameters(command="test"))

    # --- Assertions ---
    assert mock_session in group._sessions
    assert len(group.tools) == 1
    expected_tool_name = "HookServer.base_tool"
    assert expected_tool_name in group.tools
    assert group.tools[expected_tool_name] == mock_tool
    assert group._tool_to_session[expected_tool_name] == mock_session


@pytest.mark.anyio
async def test_client_session_group_disconnect_from_server():
    """Test disconnecting from a server."""
    # --- Test Setup ---
    group = ClientSessionGroup()
    server_name = "ServerToDisconnect"

    # Manually populate state using standard mocks
    mock_session1 = mock.MagicMock(spec=mcp.ClientSession)
    mock_session2 = mock.MagicMock(spec=mcp.ClientSession)
    mock_tool1 = mock.Mock(spec=types.Tool)
    mock_tool1.name = "tool1"
    mock_resource1 = mock.Mock(spec=types.Resource)
    mock_resource1.name = "res1"
    mock_prompt1 = mock.Mock(spec=types.Prompt)
    mock_prompt1.name = "prm1"
    mock_tool2 = mock.Mock(spec=types.Tool)
    mock_tool2.name = "tool2"
    mock_component_named_like_server = mock.Mock()
    mock_session = mock.Mock(spec=mcp.ClientSession)

    group._tools = {
        "tool1": mock_tool1,
        "tool2": mock_tool2,
        server_name: mock_component_named_like_server,
    }
    group._tool_to_session = {
        "tool1": mock_session1,
        "tool2": mock_session2,
        server_name: mock_session1,
    }
    group._resources = {
        "res1": mock_resource1,
        server_name: mock_component_named_like_server,
    }
    group._prompts = {
        "prm1": mock_prompt1,
        server_name: mock_component_named_like_server,
    }
    group._sessions = {
        mock_session: ClientSessionGroup._ComponentNames(
            prompts=set({"prm1"}),
            resources=set({"res1"}),
            tools=set({"tool1", "tool2"}),
        )
    }

    # --- Assertions ---
    assert mock_session in group._sessions
    assert "tool1" in group._tools
    assert "tool2" in group._tools
    assert "res1" in group._resources
    assert "prm1" in group._prompts

    # --- Test Execution ---
    await group.disconnect_from_server(mock_session)

    # --- Assertions ---
    assert mock_session not in group._sessions
    assert "tool1" not in group._tools
    assert "tool2" not in group._tools
    assert "res1" not in group._resources
    assert "prm1" not in group._prompts


@pytest.mark.anyio
async def test_client_session_group_connect_to_server_duplicate_tool_raises_error(
    mock_exit_stack: contextlib.AsyncExitStack,
):
    """Test MCPError raised when connecting a server with a dup name."""
    # --- Setup Pre-existing State ---
    group = ClientSessionGroup(exit_stack=mock_exit_stack)
    existing_tool_name = "shared_tool"
    # Manually add a tool to simulate a previous connection
    group._tools[existing_tool_name] = mock.Mock(spec=types.Tool)
    group._tools[existing_tool_name].name = existing_tool_name
    # Need a dummy session associated with the existing tool
    mock_session = mock.MagicMock(spec=mcp.ClientSession)
    group._tool_to_session[existing_tool_name] = mock_session
    group._session_exit_stacks[mock_session] = mock.Mock(spec=contextlib.AsyncExitStack)

    # --- Mock New Connection Attempt ---
    mock_server_info_new = mock.Mock(spec=types.Implementation)
    mock_server_info_new.name = "ServerWithDuplicate"
    mock_session_new = mock.AsyncMock(spec=mcp.ClientSession)

    # Configure the new session to return a tool with the *same name*
    duplicate_tool = mock.Mock(spec=types.Tool)
    duplicate_tool.name = existing_tool_name
    mock_session_new.list_tools.return_value = mock.AsyncMock(tools=[duplicate_tool])
    # Keep other lists empty for simplicity
    mock_session_new.list_resources.return_value = mock.AsyncMock(resources=[])
    mock_session_new.list_prompts.return_value = mock.AsyncMock(prompts=[])

    # --- Test Execution and Assertion ---
    with pytest.raises(MCPError) as excinfo:
        with mock.patch.object(
            group,
            "_establish_session",
            return_value=(mock_server_info_new, mock_session_new),
        ):
            await group.connect_to_server(StdioServerParameters(command="test"))

    # Assert details about the raised error
    assert excinfo.value.error.code == types.INVALID_PARAMS
    assert existing_tool_name in excinfo.value.error.message
    assert "already exist " in excinfo.value.error.message

    # Verify the duplicate tool was *not* added again (state should be unchanged)
    assert len(group._tools) == 1  # Should still only have the original
    assert group._tools[existing_tool_name] is not duplicate_tool  # Ensure it's the original mock


@pytest.mark.anyio
async def test_client_session_group_disconnect_non_existent_server():
    """Test disconnecting a server that isn't connected."""
    session = mock.Mock(spec=mcp.ClientSession)
    group = ClientSessionGroup()
    with pytest.raises(MCPError):
        await group.disconnect_from_server(session)


# TODO(Marcelo): This is horrible. We should drop this test.
@pytest.mark.anyio
@pytest.mark.parametrize(
    "server_params_instance, client_type_name, patch_target_for_client_func",
    [
        (
            StdioServerParameters(command="test_stdio_cmd"),
            "stdio",
            "mcp.client.session_group.mcp.stdio_client",
        ),
        (
            SseServerParameters(url="http://test.com/sse", timeout=10.0),
            "sse",
            "mcp.client.session_group.sse_client",
        ),  # url, headers, timeout, sse_read_timeout
        (
            StreamableHttpParameters(url="http://test.com/stream", terminate_on_close=False),
            "streamablehttp",
            "mcp.client.session_group.streamable_http_client",
        ),  # url, headers, timeout, sse_read_timeout, terminate_on_close
    ],
)
async def test_client_session_group_establish_session_parameterized(
    server_params_instance: StdioServerParameters | SseServerParameters | StreamableHttpParameters,
    client_type_name: str,  # Just for clarity or conditional logic if needed
    patch_target_for_client_func: str,
):
    with mock.patch("mcp.client.session_group.mcp.ClientSession") as mock_ClientSession_class:
        with mock.patch(patch_target_for_client_func) as mock_specific_client_func:
            mock_client_cm_instance = mock.AsyncMock(name=f"{client_type_name}ClientCM")
            mock_read_stream = mock.AsyncMock(name=f"{client_type_name}Read")
            mock_write_stream = mock.AsyncMock(name=f"{client_type_name}Write")

            # All client context managers return (read_stream, write_stream)
            mock_client_cm_instance.__aenter__.return_value = (mock_read_stream, mock_write_stream)

            mock_client_cm_instance.__aexit__ = mock.AsyncMock(return_value=None)
            mock_specific_client_func.return_value = mock_client_cm_instance

            # --- Mock mcp.ClientSession (class) ---
            # mock_ClientSession_class is already provided by the outer patch
            mock_raw_session_cm = mock.AsyncMock(name="RawSessionCM")
            mock_ClientSession_class.return_value = mock_raw_session_cm

            mock_entered_session = mock.AsyncMock(name="EnteredSessionInstance")
            mock_raw_session_cm.__aenter__.return_value = mock_entered_session
            mock_raw_session_cm.__aexit__ = mock.AsyncMock(return_value=None)

            # Mock session.initialize()
            mock_initialize_result = mock.AsyncMock(name="InitializeResult")
            mock_initialize_result.server_info = types.Implementation(name="foo", version="1")
            mock_entered_session.initialize.return_value = mock_initialize_result

            # --- Test Execution ---
            group = ClientSessionGroup()
            returned_server_info = None
            returned_session = None

            async with contextlib.AsyncExitStack() as stack:
                group._exit_stack = stack
                (
                    returned_server_info,
                    returned_session,
                ) = await group._establish_session(server_params_instance, ClientSessionParameters())

            # --- Assertions ---
            # 1. Assert the correct specific client function was called
            if client_type_name == "stdio":
                assert isinstance(server_params_instance, StdioServerParameters)
                mock_specific_client_func.assert_called_once_with(server_params_instance)
            elif client_type_name == "sse":
                assert isinstance(server_params_instance, SseServerParameters)
                mock_specific_client_func.assert_called_once_with(
                    url=server_params_instance.url,
                    headers=server_params_instance.headers,
                    timeout=server_params_instance.timeout,
                    sse_read_timeout=server_params_instance.sse_read_timeout,
                )
            elif client_type_name == "streamablehttp":  # pragma: no branch
                assert isinstance(server_params_instance, StreamableHttpParameters)
                # Verify streamable_http_client was called with url, httpx_client, and terminate_on_close
                # The http_client is created by the real create_mcp_http_client
                call_args = mock_specific_client_func.call_args
                assert call_args.kwargs["url"] == server_params_instance.url
                assert call_args.kwargs["terminate_on_close"] == server_params_instance.terminate_on_close
                assert isinstance(call_args.kwargs["http_client"], httpx.AsyncClient)

            mock_client_cm_instance.__aenter__.assert_awaited_once()

            # 2. Assert ClientSession was called correctly
            mock_ClientSession_class.assert_called_once_with(
                mock_read_stream,
                mock_write_stream,
                read_timeout_seconds=None,
                sampling_callback=None,
                elicitation_callback=None,
                list_roots_callback=None,
                logging_callback=None,
                message_handler=None,
                client_info=None,
            )
            mock_raw_session_cm.__aenter__.assert_awaited_once()
            mock_entered_session.initialize.assert_awaited_once()

            # 3. Assert returned values
            assert returned_server_info is mock_initialize_result.server_info
            assert returned_session is mock_entered_session
