import asyncio

from fastapi import FastAPI, Request
from fastapi.responses import Response
from mcp.server import Server
from mcp.server.sse import SseServerTransport
from starlette.responses import StreamingResponse

# 1. Initialize the MCP Server logic
# This is where you define your tools, resources, and prompts
mcp_server = Server("my-remote-server")


@mcp_server.list_tools()
async def handle_list_tools():
    """List available tools."""
    return [
        {
            "name": "get_weather",
            "description": "Get the current weather for a location",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "location": {"type": "string"},
                },
                "required": ["location"],
            },
        }
    ]


@mcp_server.call_tool()
async def handle_call_tool(name: str, arguments: dict):
    """Handle tool execution."""
    if name == "get_weather":
        location = arguments.get("location", "Unknown")
        # In a real app, call an actual Weather API here
        return [
            {"type": "text", "text": f"The weather in {location} is sunny and 25°C."}
        ]
    raise ValueError(f"Tool not found: {name}")


# 2. Initialize FastAPI
app = FastAPI(title="Remote MCP Server")

# 3. Create the SSE Transport layer
# This object manages the connection between the web and the MCP protocol
sse = SseServerTransport("/messages")


@app.get("/sse")
async def sse_endpoint(request: Request):
    """
    The client connects here to start the SSE stream.
    The server will push messages to the client through this connection.
    """
    async with sse.connect_sse(request.scope, request.receive, request._send) as (
        read_stream,
        write_stream,
    ):
        # We run the MCP server using the streams provided by the SSE transport
        await mcp_server.run(
            read_stream, write_stream, mcp_server.create_initialization_options()
        )


@app.post("/messages")
async def messages_endpoint(request: Request):
    """
    The client sends JSON-RPC messages (tool calls, etc.)
    via POST requests to this endpoint.
    """
    await sse.handle_post_message(request.scope, request.receive, request._send)
    return Response(status_code=202)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)