#!/usr/bin/env python3
"""14 — Connect a KOAN agent to external MCP servers.

Shows how to use any MCP-compatible server (not just KOAN's) as a tool
source for KOAN agents. This connects to:

1. A KOAN MCP server (our own tools)
2. An external MCP server (e.g., @modelcontextprotocol/server-filesystem)

Multiple MCP servers can be combined into a single agent.

Requirements:
    pip install koan[mcp,llm]
    export OPENAI_API_KEY=sk-...

    # For external filesystem MCP server (optional):
    npx -y @modelcontextprotocol/server-filesystem /tmp
"""

from __future__ import annotations

import asyncio
import sys
from typing import TYPE_CHECKING

from mcp.client.session import ClientSession
from mcp.client.stdio import StdioServerParameters, stdio_client

from koan.agents import AgentConfig, ReactAgent
from koan.llm.client import LLMClient
from koan.mcp import MCPToolSource
from koan.tools import ToolRegistry, tool

if TYPE_CHECKING:
    from koan.tools.types import ToolDefinition, ToolResult


class CombinedToolSource:
    """Merges multiple tool sources into one for agent use."""

    def __init__(self) -> None:
        self._sources: list[MCPToolSource | ToolRegistry] = []

    def add(self, source: MCPToolSource | ToolRegistry) -> None:
        self._sources.append(source)

    def definitions(self) -> list[ToolDefinition]:
        result = []
        for src in self._sources:
            result.extend(src.definitions())
        return result

    async def execute(
        self, name: str, arguments: str | dict, *, context: object = None
    ) -> ToolResult:
        for src in self._sources:
            names = {d.name for d in src.definitions()}
            if name in names:
                return await src.execute(name, arguments, context=context)
        raise KeyError(f"Tool not found across sources: '{name}'")


# --- Local tools (not MCP) ---


@tool
def get_current_time() -> str:
    """Get the current date and time."""
    from datetime import UTC, datetime

    return datetime.now(UTC).strftime("%Y-%m-%d %H:%M:%S UTC")


async def demo_koan_mcp_server() -> None:
    """Connect to our KOAN MCP server and use its tools via an agent."""
    print("=" * 60)
    print("Demo 1: KOAN agent + KOAN MCP server tools")
    print("=" * 60)

    server_params = StdioServerParameters(
        command=sys.executable,
        args=["examples/13_mcp_server.py"],
    )

    async with (
        stdio_client(server_params) as (read, write),
        ClientSession(read, write) as session,
    ):
        await session.initialize()

        # Discover MCP tools
        mcp_source = MCPToolSource(session)
        await mcp_source.discover()

        # Also add local tools
        local_tools = ToolRegistry()
        local_tools.add(get_current_time)

        # Combine both sources
        combined = CombinedToolSource()
        combined.add(mcp_source)
        combined.add(local_tools)

        all_tools = combined.definitions()
        print(f"\nAvailable tools ({len(all_tools)}):")
        for t in all_tools:
            print(f"  - {t.name}: {t.description}")

        async with LLMClient() as client:
            agent = ReactAgent(
                "multi-source-agent",
                client=client,
                config=AgentConfig(
                    provider="openai",
                    model="gpt-4o-mini",
                    system_prompt=(
                        "You have access to stock prices, document search, "
                        "employee lookup, and current time tools. "
                        "Use them to answer questions. Be concise."
                    ),
                ),
                tools=combined,
            )

            queries = [
                "What time is it and what's the stock price of MSFT?",
                "Look up employee Alice and search for 'budget' documents.",
            ]

            for query in queries:
                print(f"\nQuery: {query}")
                result = await agent.run(query)
                print(f"Answer: {result.output}")
                print(f"  ({result.tool_calls_made} tool calls)")


async def demo_external_mcp_server() -> None:
    """Connect to an external MCP server (filesystem) if available."""
    print("\n" + "=" * 60)
    print("Demo 2: KOAN agent + external MCP server (filesystem)")
    print("=" * 60)

    # Try connecting to the filesystem MCP server
    # Install with: npx -y @modelcontextprotocol/server-filesystem /tmp
    server_params = StdioServerParameters(
        command="npx",
        args=["-y", "@modelcontextprotocol/server-filesystem", "/tmp"],
    )

    try:
        async with (
            stdio_client(server_params) as (read, write),
            ClientSession(read, write) as session,
        ):
            await session.initialize()

            source = MCPToolSource(session)
            tools = await source.discover()

            print(f"\nDiscovered {len(tools)} tools from filesystem MCP server:")
            for t in tools:
                print(f"  - {t.name}: {t.description}")

            async with LLMClient() as client:
                agent = ReactAgent(
                    "fs-agent",
                    client=client,
                    config=AgentConfig(
                        provider="openai",
                        model="gpt-4o-mini",
                        system_prompt=(
                            "You have filesystem tools. Use them to explore files. "
                            "Be concise."
                        ),
                    ),
                    tools=source,
                )

                result = await agent.run("List the files in /tmp directory.")
                print("\nQuery: List files in /tmp")
                print(f"Answer: {result.output[:300]}")
                print(f"  ({result.tool_calls_made} tool calls)")

    except (OSError, FileNotFoundError):
        print("\n  [Skipped] External filesystem MCP server not available.")
        print("  Install with: npx -y @modelcontextprotocol/server-filesystem /tmp")


async def main() -> None:
    await demo_koan_mcp_server()
    await demo_external_mcp_server()


if __name__ == "__main__":
    asyncio.run(main())
