"""MCP Server — expose Agent or Team as an MCP server.

Any Runnable (Agent or Team) can be served as an MCP server,
allowing other tools and IDEs to connect to it.

Usage::

    from logos_adk.mcp import serve_as_mcp

    agent = Agent("analyst", model="claude-sonnet-4-20250514")
    serve_as_mcp(agent, port=3001)

    team = Team("content_team", agents=[...])
    serve_as_mcp(team, port=3002)
"""
from __future__ import annotations

import asyncio
from typing import Any

from logos_adk.runnable import Runnable


def _create_mcp_server(runnable: Runnable, *, name: str | None = None) -> Any:
    """Create an MCP Server instance wrapping a Runnable.

    Returns an mcp.server.Server configured with tools from the Runnable.
    """
    try:
        from mcp.server import Server
        from mcp.types import Tool as MCPTool, TextContent
    except ImportError:
        raise ImportError(
            "mcp package required for MCP server. Install with: "
            "pip install mcp"
        )

    server_name = name or f"logos-{runnable.name}"
    server = Server(server_name)

    # Register a "prompt" tool that runs the Runnable
    @server.list_tools()
    async def list_tools() -> list[MCPTool]:
        tools = [
            MCPTool(
                name="prompt",
                description=f"Send a prompt to the '{runnable.name}' agent and get a response",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "message": {
                            "type": "string",
                            "description": "The prompt message to send to the agent",
                        },
                        "session_id": {
                            "type": "string",
                            "description": "Optional session ID for conversation continuity",
                        },
                    },
                    "required": ["message"],
                },
            ),
        ]

        # If the runnable is an Agent with tools, expose them individually too
        if hasattr(runnable, '_config') and hasattr(runnable._config, 'tools'):
            for agent_tool in runnable._config.tools:
                schema = agent_tool.schema
                tools.append(MCPTool(
                    name=f"tool_{agent_tool.name}",
                    description=f"[Direct tool] {agent_tool.description}",
                    inputSchema=schema.get("parameters", {
                        "type": "object", "properties": {},
                    }),
                ))

        return tools

    @server.call_tool()
    async def call_tool(name: str, arguments: dict[str, Any] | None = None) -> list[TextContent]:
        arguments = arguments or {}

        if name == "prompt":
            message = arguments.get("message", "")
            session_id = arguments.get("session_id")
            response = await runnable.run(message, session_id=session_id)
            return [TextContent(type="text", text=response.content)]

        # Direct tool calls (tool_<name>)
        if name.startswith("tool_") and hasattr(runnable, '_config'):
            tool_name = name[5:]  # strip "tool_" prefix
            for agent_tool in runnable._config.tools:
                if agent_tool.name == tool_name:
                    result = await agent_tool(**arguments)
                    return [TextContent(type="text", text=str(result))]

        return [TextContent(type="text", text=f"Unknown tool: {name}")]

    return server


def serve_as_mcp(
    runnable: Runnable,
    *,
    port: int = 3000,
    host: str = "0.0.0.0",
    transport: str = "sse",
    name: str | None = None,
) -> None:
    """Serve a Runnable (Agent or Team) as an MCP server.

    Args:
        runnable: The Agent or Team to expose.
        port: Port to listen on (default 3000).
        host: Host to bind to (default 0.0.0.0).
        transport: Transport mode — "sse" (HTTP) or "stdio".
        name: Server name (defaults to "logos-{runnable.name}").

    Usage::

        from logos_adk import Agent
        from logos_adk.mcp import serve_as_mcp

        agent = Agent("analyst", model="claude-sonnet-4-20250514")
        serve_as_mcp(agent, port=3001)  # blocks, serves forever
    """
    server = _create_mcp_server(runnable, name=name)

    if transport == "stdio":
        _serve_stdio(server)
    elif transport == "sse":
        _serve_sse(server, host=host, port=port)
    else:
        raise ValueError(f"Unknown transport: {transport!r}. Use 'sse' or 'stdio'.")


def _serve_stdio(server: Any) -> None:
    """Serve over stdio transport."""
    from mcp.server.stdio import stdio_server

    async def _run() -> None:
        async with stdio_server() as (read_stream, write_stream):
            await server.run(read_stream, write_stream, server.create_initialization_options())

    asyncio.run(_run())


def _serve_sse(server: Any, *, host: str, port: int) -> None:
    """Serve over SSE (HTTP) transport."""
    try:
        from mcp.server.sse import SseServerTransport
        from starlette.applications import Starlette
        from starlette.routing import Route, Mount
        import uvicorn
    except ImportError:
        raise ImportError(
            "SSE server requires additional packages. Install with: "
            "pip install mcp uvicorn starlette"
        )

    sse = SseServerTransport("/messages/")

    async def handle_sse(request: Any) -> Any:
        async with sse.connect_sse(
            request.scope, request.receive, request._send
        ) as (read_stream, write_stream):
            await server.run(read_stream, write_stream, server.create_initialization_options())

    app = Starlette(
        routes=[
            Route("/sse", endpoint=handle_sse),
            Mount("/messages/", app=sse.handle_post_message),
        ],
    )

    uvicorn.run(app, host=host, port=port)


async def serve_as_mcp_async(
    runnable: Runnable,
    *,
    transport: str = "stdio",
    name: str | None = None,
) -> None:
    """Async version of serve_as_mcp for use within an event loop.

    Only supports stdio transport (SSE requires uvicorn which manages its own loop).

    Usage::

        await serve_as_mcp_async(agent, transport="stdio")
    """
    if transport != "stdio":
        raise ValueError("serve_as_mcp_async only supports 'stdio' transport")

    from mcp.server.stdio import stdio_server

    server = _create_mcp_server(runnable, name=name)
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, server.create_initialization_options())
