"""Unit tests for MCPToolProvider."""

from contextlib import asynccontextmanager
from typing import Any, AsyncContextManager, Callable
from unittest.mock import AsyncMock, patch

import pytest
from mcp import ClientSession, ListToolsResult, StdioServerParameters, Tool

from llm_agents_from_scratch.errors import (
    MCPWarning,
    MissingMCPServerParamsError,
)
from llm_agents_from_scratch.tools.mcp import MCPToolProvider


def test_mcp_tool_provider_init() -> None:
    """Tests initialization of an MCPToolProvider."""
    streamable_http_provider = MCPToolProvider(
        name="mock provider",
        streamable_http_url="https://mock-server-url.io",
    )

    assert streamable_http_provider.name == "mock provider"
    assert streamable_http_provider.stdio_params is None
    assert (
        streamable_http_provider.streamable_http_url
        == "https://mock-server-url.io"
    )

    stdio_params = StdioServerParameters(
        command="uv run",
        args=["fake.py"],
    )
    stdio_provider = MCPToolProvider(
        name="mock provider",
        stdio_params=stdio_params,
    )
    assert stdio_provider.name == "mock provider"
    assert stdio_provider.streamable_http_url is None
    assert stdio_provider.stdio_params == stdio_params


def test_mcp_tool_provider_init_raises_error() -> None:
    """Tests initialization raises error if no connection details provided."""
    with pytest.raises(
        MissingMCPServerParamsError,
        match="You must supply at least one",
    ):
        MCPToolProvider(name="invalid provider")


def test_mcp_tool_provider_init_emits_warning() -> None:
    """Tests init emits warning if both connection details provided."""
    with pytest.warns(
        MCPWarning,
        match="Both `stdio_params` and `streamable_http_url`",
    ):
        stdio_params = StdioServerParameters(
            command="uv run",
            args=["fake.py"],
        )
        MCPToolProvider(
            name="mock provider",
            stdio_params=stdio_params,
            streamable_http_url="https://mock-server-url.io",
        )


@pytest.fixture()
def mock_client_session() -> Callable[..., AsyncContextManager[AsyncMock]]:
    """Mock ClientSession."""

    @asynccontextmanager
    async def async_client_session(*args, **kwargs):
        client_session = AsyncMock(spec=ClientSession)
        mock_list_tools = AsyncMock()
        mock_list_tools.return_value = ListToolsResult(
            tools=[
                Tool(
                    name="mock_tool",
                    description="mock_desc",
                    inputSchema={"param1": {"type": "number"}},
                ),
            ],
        )
        client_session.list_tools = mock_list_tools
        yield client_session

    return async_client_session


@pytest.mark.asyncio
@patch("llm_agents_from_scratch.tools.mcp.provider.stdio_client")
@patch("llm_agents_from_scratch.tools.mcp.provider.ClientSession")
async def test_session_creation(
    mock_client_session_cls: AsyncMock,
    mock_stdio_client: AsyncMock,
    mock_stdio_client_transport: AsyncContextManager[Any],
    mock_client_session: Callable[..., AsyncContextManager[AsyncMock]],
) -> None:
    """Tests creation of sessions."""
    # Set up the mock to return the async context manager
    mock_stdio_client.side_effect = mock_stdio_client_transport
    mock_client_session_cls.side_effect = mock_client_session

    stdio_params = StdioServerParameters(
        command="uv run",
        args=["fake.py"],
    )
    stdio_provider = MCPToolProvider(
        name="mock provider",
        stdio_params=stdio_params,
    )

    await stdio_provider.session()

    mock_stdio_client.assert_called_once()
    mock_client_session_cls.assert_called_once()
    assert stdio_provider._session_ready.is_set()


@pytest.mark.asyncio
@patch("llm_agents_from_scratch.tools.mcp.provider.streamablehttp_client")
@patch("llm_agents_from_scratch.tools.mcp.provider.ClientSession")
async def test_session_creation_streamable_http(
    mock_client_session_cls: AsyncMock,
    mock_streamablehttp_client: AsyncMock,
    mock_streamable_http_client_transport: Callable[
        ...,
        AsyncContextManager[Any],
    ],
    mock_client_session: Callable[..., AsyncContextManager[AsyncMock]],
) -> None:
    """Tests creation of sessions."""
    # Set up the mock to return the async context manager
    mock_streamablehttp_client.side_effect = (
        mock_streamable_http_client_transport
    )
    mock_client_session_cls.side_effect = mock_client_session

    streamablehttp_provider = MCPToolProvider(
        name="mock provider",
        streamable_http_url="http://mock-url.io",
    )

    await streamablehttp_provider.session()

    mock_streamablehttp_client.assert_called_once()
    mock_client_session_cls.assert_called_once()
    assert streamablehttp_provider._session_ready.is_set()


@pytest.mark.asyncio
@patch("llm_agents_from_scratch.tools.mcp.provider.stdio_client")
@patch("llm_agents_from_scratch.tools.mcp.provider.ClientSession")
async def test_session_creation_raises_error(
    mock_client_session_cls: AsyncMock,
    mock_stdio_client: AsyncMock,
    mock_stdio_client_transport: AsyncContextManager[Any],
    mock_client_session: Callable[..., AsyncContextManager[AsyncMock]],
) -> None:
    """Tests creation of sessions."""
    # Set up the mock to return the async context manager
    mock_stdio_client.side_effect = mock_stdio_client_transport
    mock_client_session_cls.side_effect = FileNotFoundError()

    stdio_params = StdioServerParameters(
        command="uv run",
        args=["fake.py"],
    )
    stdio_provider = MCPToolProvider(
        name="mock provider",
        stdio_params=stdio_params,
    )

    with pytest.raises(FileNotFoundError):
        await stdio_provider.session()


@pytest.mark.asyncio
@patch("llm_agents_from_scratch.tools.mcp.provider.stdio_client")
@patch("llm_agents_from_scratch.tools.mcp.provider.ClientSession")
async def test_list_tools_stdio_client(
    mock_client_session_cls: AsyncMock,
    mock_stdio_client: AsyncMock,
    mock_stdio_client_transport: AsyncContextManager[Any],
    mock_client_session: Callable[..., AsyncContextManager[AsyncMock]],
) -> None:
    """Tests creation of sessions."""
    # Set up the mock to return the async context manager
    mock_stdio_client.side_effect = mock_stdio_client_transport
    mock_client_session_cls.side_effect = mock_client_session

    stdio_params = StdioServerParameters(
        command="uv run",
        args=["fake.py"],
    )
    stdio_provider = MCPToolProvider(
        name="mock_provider",
        stdio_params=stdio_params,
    )

    mcp_tools = await stdio_provider.get_tools()

    assert len(mcp_tools) == 1
    assert mcp_tools[0].description == "mock_desc"
    assert mcp_tools[0].name == "mcp__mock_provider__mock_tool"
    assert mcp_tools[0].parameters_json_schema == {"param1": {"type": "number"}}
    assert mcp_tools[0].additional_annotations is None


@pytest.mark.asyncio
@patch("llm_agents_from_scratch.tools.mcp.provider.stdio_client")
@patch("llm_agents_from_scratch.tools.mcp.provider.ClientSession")
async def test_close(
    mock_client_session_cls: AsyncMock,
    mock_stdio_client: AsyncMock,
    mock_stdio_client_transport: AsyncContextManager[Any],
    mock_client_session: Callable[..., AsyncContextManager[AsyncMock]],
) -> None:
    """Test closing of session."""
    # Set up the mock to return the async context manager
    mock_stdio_client.side_effect = mock_stdio_client_transport
    mock_client_session_cls.side_effect = mock_client_session

    stdio_params = StdioServerParameters(
        command="uv run",
        args=["fake.py"],
    )
    stdio_provider = MCPToolProvider(
        name="mock_provider",
        stdio_params=stdio_params,
    )

    await stdio_provider.session()
    assert stdio_provider._session_ready.is_set()
    assert not stdio_provider._shutdown_event.is_set()

    # act
    await stdio_provider.close()
    assert not stdio_provider._shutdown_event.is_set()
    assert stdio_provider._session is None
    assert stdio_provider._session_task is None
