import asyncio
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import TypedDict

from fastmcp import Client, FastMCP
from fastmcp.server.context import Context


class AppState(TypedDict):
    message: str
    counter: int


@asynccontextmanager
async def lifespan(app: FastMCP[AppState]) -> AsyncIterator[AppState]:
    print("Lifespan: starting", flush=True)
    try:
        yield {"message": "Hello from lifespan!", "counter": 42}
    finally:
        print("Lifespan: shutdown", flush=True)


mcp = FastMCP[AppState]("test-server", lifespan=lifespan)


@mcp.tool
async def get_message(ctx: Context) -> str:
    """Get the message from lifespan state."""
    # PROBLEM: _lifespan_result is private, no public API
    state = ctx.fastmcp._lifespan_result
    if state is None:
        raise RuntimeError("Lifespan state not available")
    return f"Message: {state['message']}, Counter: {state['counter']}"


async def main():
    async with Client(mcp) as client:
        result = await client.call_tool("get_message", {})
        print(f"Result: {result.content[0].text}")

        # Verify the state was accessible
        assert "Hello from lifespan!" in result.content[0].text
        assert "42" in result.content[0].text
        print("SUCCESS: Lifespan state is accessible via _lifespan_result")


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

# Lifespan: starting
# Result: Message: Hello from lifespan!, Counter: 42
# SUCCESS: Lifespan state is accessible via _lifespan_result
# Lifespan: shutdown