from mcp.server.fastmcp import FastMCP, Context
from mcp.server.session import ServerSession
from mcp.types import SamplingMessage, TextContent


# MCPサーバの準備
mcp = FastMCP("sampling_serve")


@mcp.tool()
async def generate_poem(topic: str, ctx: Context[ServerSession, None]) -> str:
    """サンプリングで詩を作成"""
    prompt = f"{topic}の詩を作成して。"
    result = await ctx.session.create_message(
        messages=[
            SamplingMessage(
                role="user",
                content=TextContent(type="text", text=prompt),
            )
        ],
        max_tokens=1000,
    )

    if result.content.type == "text":
        return result.content.text
    return str(result.content)


if __name__ == "__main__":
    """メイン"""
    mcp.run(transport="stdio")