from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass

from mcp.server.fastmcp import Context, FastMCP
from mcp.server.session import ServerSession


class Database:  # Replace with your actual DB type
    @classmethod
    async def connect(cls):  # pragma: no cover
        return cls()

    async def disconnect(self):  # pragma: no cover
        pass

    def query(self):  # pragma: no cover
        return "Hello, World!"


# Create a named server
mcp = FastMCP("My App")


@dataclass
class AppContext:
    db: Database


@asynccontextmanager
async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:  # pragma: no cover
    """Manage application lifecycle with type-safe context"""
    # Initialize on startup
    db = await Database.connect()
    try:
        yield AppContext(db=db)
    finally:
        # Cleanup on shutdown
        await db.disconnect()


# Pass lifespan to server
mcp = FastMCP("My App", lifespan=app_lifespan)


# Access type-safe lifespan context in tools
@mcp.tool()
def query_db(ctx: Context[ServerSession, AppContext]) -> str:  # pragma: no cover
    """Tool that uses initialized resources"""
    db = ctx.request_context.lifespan_context.db
    return db.query()
