#!/usr/bin/env python3
"""15 — Supervisor with per-agent MCP servers.

Each sub-agent connects to its own dedicated MCP server, giving it a
specialized toolset. The Supervisor decomposes compound queries and
routes sub-tasks to the agent with the right tools.

Architecture:
    ┌─────────────────────────────────────────────────┐
    │                  Supervisor                      │
    │  (decompose → delegate → synthesize)            │
    │                                                  │
    │  ┌───────────────┐    ┌───────────────┐         │
    │  │ finance-agent  │    │   hr-agent    │         │
    │  │  (ReactAgent)  │    │ (ReactAgent)  │         │
    │  └───────┬───────┘    └──────┬────────┘         │
    └──────────┼───────────────────┼──────────────────┘
               │                   │
    ┌──────────▼───────┐  ┌───────▼─────────┐
    │  Finance MCP     │  │    HR MCP       │
    │  Server (stdio)  │  │  Server (stdio) │
    │  - stock prices  │  │  - employees    │
    │  - portfolios    │  │  - teams        │
    │  - currency      │  │  - PTO balance  │
    └──────────────────┘  └─────────────────┘

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

from __future__ import annotations

import asyncio
import sys
from contextlib import AsyncExitStack

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

from koan.agents import (
    AgentComplete,
    AgentConfig,
    AgentError,
    AgentStart,
    ReactAgent,
    RuleRouter,
    Supervisor,
    TextDelta,
    ToolCallEnd,
    ToolCallStart,
)
from koan.llm.client import LLMClient
from koan.llm.config import LLMConfig
from koan.mcp import MCPToolSource


async def connect_mcp_agent(
    stack: AsyncExitStack,
    name: str,
    server_script: str,
    system_prompt: str,
    client: LLMClient,
) -> tuple[ReactAgent, MCPToolSource]:
    """Launch an MCP server subprocess and create an agent connected to it.

    Uses AsyncExitStack so the server stays alive for the whole session.
    """
    server_params = StdioServerParameters(
        command=sys.executable,
        args=[server_script],
    )

    transport = await stack.enter_async_context(stdio_client(server_params))
    read_stream, write_stream = transport
    session = await stack.enter_async_context(ClientSession(read_stream, write_stream))
    await session.initialize()

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

    agent = ReactAgent(
        name,
        client=client,
        config=AgentConfig(
            provider="openai",
            model="gpt-4o",
            system_prompt=system_prompt,
        ),
        tools=source,
    )

    return agent, source


async def demo_non_streaming() -> None:
    """Supervisor.run() — final results only."""
    print("=" * 60)
    print("MODE: Non-Streaming (supervisor.run)")
    print("=" * 60)

    async with AsyncExitStack() as stack:
        client = await stack.enter_async_context(LLMClient())

        finance_agent, finance_src = await connect_mcp_agent(
            stack,
            name="finance",
            server_script="examples/15_mcp_finance_server.py",
            system_prompt=(
                "You are a finance specialist. Use your tools to answer "
                "questions about stocks, portfolios, and currency. Be concise."
            ),
            client=client,
        )

        hr_agent, hr_src = await connect_mcp_agent(
            stack,
            name="hr",
            server_script="examples/15_mcp_hr_server.py",
            system_prompt=(
                "You are an HR specialist. Use your tools to answer questions "
                "about employees, teams, and PTO. Be concise."
            ),
            client=client,
        )

        print(f"\nFinance tools: {[t.name for t in finance_src.definitions()]}")
        print(f"HR tools:      {[t.name for t in hr_src.definitions()]}")

        supervisor = Supervisor(
            agents={"finance": finance_agent, "hr": hr_agent},
            router=RuleRouter(
                rules={
                    "stock": "finance",
                    "price": "finance",
                    "portfolio": "finance",
                    "currency": "finance",
                    "convert": "finance",
                    "employee": "hr",
                    "team": "hr",
                    "pto": "hr",
                    "who": "hr",
                },
            ),
            llm_client=client,
            llm_config=LLMConfig(provider="openai", model="gpt-4o"),
            fallback="finance",
        )

        queries = [
            "What's the stock price of AAPL and how many PTO days does Alice have?",
            "Show me Bob's portfolio and look up which team Carol is on.",
            "Convert 100 USD to EUR.",
        ]

        for query in queries:
            print(f"\nQuery: {query}")
            print("-" * 40)
            result = await supervisor.run(query)
            print(f"Agent: {result.agent_name}")
            print(f"Output: {result.output}")
            print(f"Tool calls: {result.tool_calls_made}")


async def demo_streaming() -> None:
    """Supervisor.run_stream() — see sub-agent work live."""
    print("\n" + "=" * 60)
    print("MODE: Streaming (supervisor.run_stream)")
    print("=" * 60)

    async with AsyncExitStack() as stack:
        client = await stack.enter_async_context(LLMClient())

        finance_agent, _ = await connect_mcp_agent(
            stack,
            name="finance",
            server_script="examples/15_mcp_finance_server.py",
            system_prompt=(
                "You are a finance specialist. Use your tools to answer "
                "questions about stocks, portfolios, and currency. Be concise."
            ),
            client=client,
        )

        hr_agent, _ = await connect_mcp_agent(
            stack,
            name="hr",
            server_script="examples/15_mcp_hr_server.py",
            system_prompt=(
                "You are an HR specialist. Use your tools to answer questions "
                "about employees, teams, and PTO. Be concise."
            ),
            client=client,
        )

        supervisor = Supervisor(
            agents={"finance": finance_agent, "hr": hr_agent},
            router=RuleRouter(
                rules={
                    "stock": "finance",
                    "price": "finance",
                    "portfolio": "finance",
                    "currency": "finance",
                    "convert": "finance",
                    "employee": "hr",
                    "team": "hr",
                    "pto": "hr",
                    "who": "hr",
                },
            ),
            llm_client=client,
            llm_config=LLMConfig(provider="openai", model="gpt-4o"),
            fallback="finance",
        )

        query = (
            "What's the stock price of MSFT, who are the members of Team Alpha, "
            "and how many PTO days does Bob have?"
        )

        print(f"\nQuery: {query}")
        print("-" * 40)

        async for event in supervisor.run_stream(query):
            if isinstance(event, AgentStart):
                print(f"\n  [Delegated to: {event.agent_name}] {event.input_text}")
            elif isinstance(event, TextDelta):
                print(event.text, end="", flush=True)
            elif isinstance(event, ToolCallStart):
                print(f"\n    -> {event.tool_name}({event.arguments})")
            elif isinstance(event, ToolCallEnd):
                print(f"    <- {event.result}")
            elif isinstance(event, AgentComplete):
                if event.metadata.get("synthesized"):
                    print("\n\n  [Supervisor synthesized answer]")
                    print(f"  {event.output}")
                else:
                    print(
                        f"\n  [Done: {event.iterations} iters, "
                        f"{event.tool_calls_made} tool calls]"
                    )
            elif isinstance(event, AgentError):
                print(f"\n  [Error: {event.error}]")

        print()


async def main() -> None:
    await demo_non_streaming()
    await demo_streaming()


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