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

from mcp.server.mcpserver import Context, MCPServer


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 = MCPServer("My App")


@dataclass
class AppContext:
    db: Database


@asynccontextmanager
async def app_lifespan(server: MCPServer) -> 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 = MCPServer("My App", lifespan=app_lifespan)


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