"""Tests for KoanMCPServer and MCPToolSource using in-memory transport."""

from __future__ import annotations

import pytest
from mcp.shared.memory import create_connected_server_and_client_session

from koan.mcp.client import MCPToolSource
from koan.mcp.server import KoanMCPServer
from koan.tools.decorator import tool
from koan.tools.registry import ToolRegistry


@tool
def greet(name: str) -> str:
    """Greet someone by name."""
    return f"Hello, {name}!"


@tool
def add(a: int, b: int) -> int:
    """Add two numbers."""
    return a + b


@tool
async def async_multiply(x: int, y: int) -> int:
    """Multiply two numbers asynchronously."""
    return x * y


def _make_registry() -> ToolRegistry:
    registry = ToolRegistry()
    registry.add(greet)
    registry.add(add)
    registry.add(async_multiply)
    return registry


class TestKoanMCPServer:
    async def test_tools_registered(self) -> None:
        """KOAN tools appear in MCP list_tools."""
        registry = _make_registry()
        server = KoanMCPServer("test-server", registry=registry)

        async with create_connected_server_and_client_session(server.mcp) as session:
            await session.initialize()
            result = await session.list_tools()

            tool_names = {t.name for t in result.tools}
            assert "greet" in tool_names
            assert "add" in tool_names
            assert "async_multiply" in tool_names

    async def test_call_sync_tool(self) -> None:
        """MCP call_tool executes a sync KOAN tool."""
        registry = _make_registry()
        server = KoanMCPServer("test-server", registry=registry)

        async with create_connected_server_and_client_session(server.mcp) as session:
            await session.initialize()
            result = await session.call_tool("greet", arguments={"name": "World"})

            assert not result.isError
            text = result.content[0].text  # type: ignore[union-attr]
            assert "Hello, World!" in text

    async def test_call_async_tool(self) -> None:
        """MCP call_tool executes an async KOAN tool."""
        registry = _make_registry()
        server = KoanMCPServer("test-server", registry=registry)

        async with create_connected_server_and_client_session(server.mcp) as session:
            await session.initialize()
            result = await session.call_tool("async_multiply", arguments={"x": 6, "y": 7})

            assert not result.isError
            text = result.content[0].text  # type: ignore[union-attr]
            assert "42" in text

    async def test_call_tool_with_int_result(self) -> None:
        """MCP call_tool handles integer return values."""
        registry = _make_registry()
        server = KoanMCPServer("test-server", registry=registry)

        async with create_connected_server_and_client_session(server.mcp) as session:
            await session.initialize()
            result = await session.call_tool("add", arguments={"a": 3, "b": 4})

            assert not result.isError
            text = result.content[0].text  # type: ignore[union-attr]
            assert "7" in text


class TestMCPToolSource:
    async def test_discover_tools(self) -> None:
        """MCPToolSource discovers tools from the server."""
        registry = _make_registry()
        server = KoanMCPServer("test-server", registry=registry)

        async with create_connected_server_and_client_session(server.mcp) as session:
            await session.initialize()
            source = MCPToolSource(session)
            tools = await source.discover()

            assert len(tools) == 3
            assert "greet" in source
            assert "add" in source
            assert source.get("greet") is not None

    async def test_execute_via_source(self) -> None:
        """MCPToolSource can execute a discovered tool."""
        registry = _make_registry()
        server = KoanMCPServer("test-server", registry=registry)

        async with create_connected_server_and_client_session(server.mcp) as session:
            await session.initialize()
            source = MCPToolSource(session)
            await source.discover()

            result = await source.execute("greet", {"name": "KOAN"})
            assert not result.is_error
            assert "Hello, KOAN!" in str(result.output)

    async def test_execute_unknown_tool_raises(self) -> None:
        """Executing an unknown tool raises KeyError."""
        registry = _make_registry()
        server = KoanMCPServer("test-server", registry=registry)

        async with create_connected_server_and_client_session(server.mcp) as session:
            await session.initialize()
            source = MCPToolSource(session)
            await source.discover()

            with pytest.raises(KeyError, match="not_a_tool"):
                await source.execute("not_a_tool", {})

    async def test_definitions_returns_tool_definitions(self) -> None:
        """definitions() returns ToolDefinition objects."""
        registry = _make_registry()
        server = KoanMCPServer("test-server", registry=registry)

        async with create_connected_server_and_client_session(server.mcp) as session:
            await session.initialize()
            source = MCPToolSource(session)
            await source.discover()

            defns = source.definitions()
            assert len(defns) == 3
            names = {d.name for d in defns}
            assert names == {"greet", "add", "async_multiply"}

    async def test_len_and_contains(self) -> None:
        """__len__ and __contains__ work correctly."""
        registry = _make_registry()
        server = KoanMCPServer("test-server", registry=registry)

        async with create_connected_server_and_client_session(server.mcp) as session:
            await session.initialize()
            source = MCPToolSource(session)

            assert len(source) == 0
            await source.discover()
            assert len(source) == 3
            assert "greet" in source
            assert "nonexistent" not in source
