from mcp.server.fastmcp import FastMCP
from starlette.applications import Starlette
from starlette.routing import Mount
import contextlib
import uvicorn

# Creates a server named "Arithmetic"
mcp = FastMCP("Arithmetic")

@mcp.tool()
def add(a: int, b: int) -> int:
    """Add two numbers"""
    return a + b

@mcp.tool()
def multiply(a: int, b: int) -> int:
    """Multiply two numbers"""
    return a * b

@mcp.tool()
def minus(a: int, b: int) -> int:
    """Subtract two numbers (a - b)"""
    return a - b

@mcp.tool()
def divide(a: int, b: int) -> float:
    """Divide two numbers (a / b). Returns a float. Raises ValueError on division by zero."""
    if b == 0:
        raise ValueError("Division by zero is not allowed.")
    return a / b

@contextlib.asynccontextmanager
async def lifespan(app: Starlette):
    async with contextlib.AsyncExitStack() as stack:
        await stack.enter_async_context(mcp.session_manager.run())
        yield

app = Starlette(
    routes=[Mount("/mcp", mcp.streamable_http_app())],
    lifespan=lifespan,
)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8001)