"""
Nexus Demo — MCP (Model Context Protocol) Integration

Demonstrates connecting agents to external tools via the MCP protocol:
  1. Local MCP server  — filesystem access via stdio transport
  2. HTTP MCP server   — remote tool server via HTTP transport
  3. Mixed mode        — agent with both native tools and MCP tools

Prerequisites:
    pip install 'nexus-agents[mcp]'

    # For the filesystem demo (optional, needs Node.js):
    npx @modelcontextprotocol/server-filesystem /tmp

Usage:
    python examples/demo_mcp.py
    python examples/demo_mcp.py local    # stdio MCP server
    python examples/demo_mcp.py http     # HTTP MCP server
    python examples/demo_mcp.py mixed    # native + MCP tools
"""

import asyncio
import json
import sys

from _config import get_model_config

from nexus import Team
from nexus.core.agent import AgentTool
from nexus.orchestration import TopologyMode
from nexus.protocol.mcp import MCPToolProvider


# ================================================================
# Demo 1: Local MCP server (stdio transport)
# ================================================================


async def demo_local_mcp():
    """Connect to a local MCP server via stdio and use its tools."""
    print("=" * 60)
    print("Demo 1: Local MCP Server (stdio transport)")
    print("=" * 60)

    # Check if npx is available
    import shutil

    if not shutil.which("npx"):
        print(
            "\nSkipped: 'npx' not found. Install Node.js to run this demo.\n"
            "  brew install node   # macOS\n"
            "  apt install nodejs  # Linux"
        )
        return

    provider = MCPToolProvider(
        server_command=["npx", "@modelcontextprotocol/server-filesystem", "/tmp"],
        name="filesystem",
    )

    try:
        await provider.connect()
        tools = provider.get_tools()
        print(f"\nDiscovered {len(tools)} tools from MCP server:")
        for t in tools:
            print(f"  - {t.name}: {t.description[:60]}")

        if tools:
            team = Team(**get_model_config(), mode=TopologyMode.ROUTER, verbose=False)
            team.add_agent(
                "file_agent",
                instructions=(
                    "You are a file system assistant. Use the provided tools to "
                    "read, write, and manage files. Be helpful and concise."
                ),
                tools=tools,
            )

            print("\nUser: List the files in /tmp")
            result = await team.run("List the files in /tmp")
            print(f"Assistant: {result[:300]}")
            await team.close()

    except Exception as exc:
        print(f"\nMCP connection failed (this is expected without the server): {exc}")
        print("To install the filesystem server: npm install -g @modelcontextprotocol/server-filesystem")

    finally:
        await provider.disconnect()


# ================================================================
# Demo 2: HTTP MCP server (simulated)
# ================================================================


async def demo_http_mcp():
    """Demonstrate MCP over HTTP transport with a simulated server."""
    print("\n" + "=" * 60)
    print("Demo 2: HTTP MCP Server (simulated)")
    print("=" * 60)

    # Since we may not have a real MCP HTTP server running,
    # we demonstrate the pattern by creating tools manually
    # that follow the MCP protocol interface.

    print("\nNote: This demo shows the MCP HTTP integration pattern.")
    print("In production, MCPToolProvider connects to a real HTTP endpoint.\n")

    # Simulate what MCPToolProvider.get_tools() would return
    async def mcp_weather_tool(city: str) -> str:
        """Get current weather for a city (simulated MCP tool).

        Args:
            city: City name to get weather for
        """
        # In reality, this calls the MCP server via JSON-RPC
        weather_data = {
            "Beijing": {"temp": "5C", "condition": "Cloudy"},
            "Tokyo": {"temp": "12C", "condition": "Sunny"},
            "New York": {"temp": "3C", "condition": "Snow"},
            "London": {"temp": "8C", "condition": "Rainy"},
        }
        data = weather_data.get(
            city, {"temp": "N/A", "condition": "Unknown city"}
        )
        return json.dumps({"city": city, **data}, ensure_ascii=False)

    async def mcp_translate_tool(text: str, target_language: str) -> str:
        """Translate text to another language (simulated MCP tool).

        Args:
            text: Text to translate
            target_language: Target language code (e.g. "zh", "ja", "en")
        """
        return f"[MCP Translation] '{text}' -> ({target_language}) [simulated result]"

    team = Team(**get_model_config(), mode=TopologyMode.ROUTER, verbose=False)

    team.add_agent(
        "travel_assistant",
        instructions=(
            "You are a travel assistant with access to weather and translation tools. "
            "Help users plan their trips. Be concise."
        ),
        tools=[mcp_weather_tool, mcp_translate_tool],
    )

    print("User: What is the weather like in Beijing and Tokyo?")
    result = await team.run("What is the weather like in Beijing and Tokyo?")
    print(f"Assistant: {result[:400]}")

    await team.close()


# ================================================================
# Demo 3: Mixed mode — native tools + MCP tools on same agent
# ================================================================


async def demo_mixed():
    """Agent with both native Python tools and MCP-provided tools."""
    print("\n" + "=" * 60)
    print("Demo 3: Mixed Mode (native + MCP tools)")
    print("=" * 60)

    import math

    # Native tool
    def calculate(expression: str) -> str:
        """Evaluate a math expression.

        Args:
            expression: A Python math expression
        """
        allowed = {"__builtins__": {}, "math": math, "abs": abs, "round": round}
        try:
            return str(eval(expression, allowed))  # noqa: S307
        except Exception as exc:
            return f"Error: {exc}"

    # Simulated MCP tool (in production, from MCPToolProvider)
    async def mcp_unit_converter(value: float, from_unit: str, to_unit: str) -> str:
        """Convert between units (simulated MCP tool).

        Args:
            value: The numeric value to convert
            from_unit: Source unit (e.g. "km", "miles", "celsius", "fahrenheit")
            to_unit: Target unit
        """
        conversions = {
            ("km", "miles"): lambda v: v * 0.621371,
            ("miles", "km"): lambda v: v * 1.60934,
            ("celsius", "fahrenheit"): lambda v: v * 9 / 5 + 32,
            ("fahrenheit", "celsius"): lambda v: (v - 32) * 5 / 9,
            ("kg", "pounds"): lambda v: v * 2.20462,
            ("pounds", "kg"): lambda v: v * 0.453592,
        }
        key = (from_unit.lower(), to_unit.lower())
        if key in conversions:
            result = conversions[key](value)
            return f"{value} {from_unit} = {result:.2f} {to_unit}"
        return f"Unsupported conversion: {from_unit} -> {to_unit}"

    team = Team(**get_model_config(), mode=TopologyMode.ROUTER, verbose=False)

    team.add_agent(
        "smart_assistant",
        instructions=(
            "You are a smart assistant with a calculator (native tool) "
            "and a unit converter (MCP tool). Use the right tool for each question."
        ),
        tools=[calculate, mcp_unit_converter],
    )

    prompts = [
        "What is 2^16?",
        "Convert 100 km to miles",
        "Calculate math.pi * 10**2, then convert 314.159 kg to pounds",
    ]
    for prompt in prompts:
        print(f"\nUser: {prompt}")
        result = await team.run(prompt)
        print(f"Assistant: {result[:300]}")

    await team.close()


# ================================================================
# Main
# ================================================================


async def main():
    demos = {"local": demo_local_mcp, "http": demo_http_mcp, "mixed": demo_mixed}
    choice = sys.argv[1] if len(sys.argv) > 1 else "all"

    if choice == "all":
        for fn in demos.values():
            await fn()
    elif choice in demos:
        await demos[choice]()
    else:
        print(f"Usage: python demo_mcp.py [{'|'.join(demos.keys())}|all]")


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