"""MCPToolSource — discover tools from an MCP server and use them in KOAN."""

from __future__ import annotations

import json
from typing import TYPE_CHECKING, Any

from koan.mcp.convert import mcp_tool_to_koan
from koan.tools.types import ToolDefinition, ToolResult

if TYPE_CHECKING:
    from mcp.client.session import ClientSession


class MCPToolSource:
    """Connects to an MCP server and exposes its tools for KOAN agents.

    Discovered tools can be listed, and called via ``execute()``.

    Usage::

        async with stdio_client(params) as (read, write):
            async with ClientSession(read, write) as session:
                await session.initialize()
                source = MCPToolSource(session)
                await source.discover()

                # List available tools
                for defn in source.definitions():
                    print(defn.name, defn.description)

                # Execute a tool
                result = await source.execute("tool_name", {"arg": "value"})
    """

    def __init__(self, session: ClientSession) -> None:
        self._session = session
        self._tools: dict[str, ToolDefinition] = {}

    async def discover(self) -> list[ToolDefinition]:
        """Discover tools from the MCP server.

        Returns:
            List of ToolDefinition objects for available tools.
        """
        result = await self._session.list_tools()
        self._tools.clear()

        for tool in result.tools:
            info = mcp_tool_to_koan(
                name=tool.name,
                description=tool.description,
                input_schema=tool.inputSchema,
            )
            defn = ToolDefinition(
                name=info["name"],
                description=info["description"],
                parameters=info["parameters"],
            )
            self._tools[defn.name] = defn

        return list(self._tools.values())

    def definitions(self) -> list[ToolDefinition]:
        """Return discovered tool definitions."""
        return list(self._tools.values())

    def get(self, name: str) -> ToolDefinition | None:
        """Look up a discovered tool by name."""
        return self._tools.get(name)

    async def execute(
        self,
        name: str,
        arguments: dict[str, Any] | str | None = None,
        *,
        context: Any | None = None,
    ) -> ToolResult:
        """Execute a tool on the remote MCP server.

        Args:
            name: Tool name.
            arguments: Tool arguments dict.

        Returns:
            ToolResult with the output or error.
        """
        if name not in self._tools:
            raise KeyError(f"Tool not found: '{name}'")

        args = json.loads(arguments) if isinstance(arguments, str) else arguments or {}

        result = await self._session.call_tool(name, arguments=args)

        # Extract text from content blocks
        output_parts: list[str] = []
        for content in result.content:
            text = getattr(content, "text", None)
            if text is not None:
                output_parts.append(text)

        output_text = "\n".join(output_parts)

        if result.isError:
            return ToolResult(tool_name=name, error=output_text)

        return ToolResult(tool_name=name, output=output_text)

    def __len__(self) -> int:
        return len(self._tools)

    def __contains__(self, name: str) -> bool:
        return name in self._tools
