from __future__ import annotations

from typing import TYPE_CHECKING, Any, override

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

from fastmcp import Context, FastMCP
from fastmcp.server.middleware import Middleware, MiddlewareContext
from key_value.aio.stores.memory import MemoryStore

fix_bug = False
shared_store = MemoryStore() if fix_bug else None


class AddContextMiddleware(Middleware):
    """Middleware that adds context info to each request."""

    @override
    async def on_request(
        self,
        context: MiddlewareContext,
        call_next: Callable[..., Awaitable[Any]],
    ) -> Any:
        """Set context info state before passing to next handler.

        Returns:
            The response from the next handler.

        """
        if context.fastmcp_context:
            context.fastmcp_context.set_state("info", "I'm some context")
        return await call_next(context)


mcp_parent = FastMCP(
    "Parent",
    middleware=[AddContextMiddleware()],
    session_state_store=shared_store,
)

mcp_child = FastMCP("Child", session_state_store=shared_store)
mcp_parent.mount(mcp_child)


@mcp_child.tool
def print_context(ctx: Context) -> str:
    """Print the context info state.

    Returns:
        The info state value, or None because states are not inherited by default.

    """
    return ctx.get_state("info")


def main() -> None:
    """Run the parent MCP server."""
    mcp_parent.run()