#!/usr/bin/env python3
"""13 — MCP server + client: expose tools via MCP, then use them in a KOAN agent.

This example demonstrates the full MCP loop:
1. Defines a KOAN MCP server script with custom tools
2. Launches the server as a subprocess
3. Connects a KOAN agent to the server via MCPToolSource
4. The agent discovers and calls tools through MCP

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

from __future__ import annotations

import asyncio
import sys

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


async def main() -> None:
    # Point at the MCP server script — it will be launched as a subprocess
    server_params = StdioServerParameters(
        command=sys.executable,
        args=["examples/13_mcp_server.py"],
    )

    print("Connecting to KOAN MCP server...")

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

        # Discover tools from the MCP server
        source = MCPToolSource(session)
        tools = await source.discover()

        print(f"Discovered {len(tools)} tools:")
        for t in tools:
            print(f"  - {t.name}: {t.description}")

        # Use MCPToolSource as the tool provider for a KOAN agent
        async with LLMClient() as client:
            agent = ReactAgent(
                "mcp-agent",
                client=client,
                config=AgentConfig(
                    provider="openai",
                    model="gpt-4o-mini",
                    system_prompt=(
                        "You have access to tools via MCP. "
                        "Use them to answer questions. Be concise."
                    ),
                ),
                tools=source,
            )

            queries = [
                "What is the stock price of AAPL?",
                "Search documents for 'quarterly report'",
            ]

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


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