from pydantic import BaseModel, Field
from mcp.server.fastmcp import Context, FastMCP
from mcp.server.session import ServerSession


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


class BookingPreferences(BaseModel):
    """ ユーザーから情報収集するためのスキーマ """
    checkAlternative: bool = Field(description="別の日付を確認しますか？")
    alternativeDate: str = Field(
        default="2025-12-26",
        description="代替日付 (YYYY-MM-DD)",
    )


@mcp.tool()
async def book_table(date: str, time: str, party_size: int, ctx: Context[ServerSession, None]) -> str:
    """日付の空き状況を確認してテーブルを予約"""
    # 日付が空いているか確認
    if date == "2025-12-25":
        # 日付が利用できません - ユーザーに代替日を尋ねてください
        result = await ctx.elicit(
            message=(f"{date} の {party_size}名用テーブルは満席です。別の日付を確認しますか？"),
            schema=BookingPreferences,
        )

        if result.action == "accept" and result.data:
            if result.data.checkAlternative:
                return f"[予約成功] {result.data.alternativeDate} に予約しました"
            return "[予約中止] 別の日付は選択されませんでした"
        return "[予約中止] ユーザーによりキャンセルされました"

    # 利用可能日
    return f"[予約成功] {date} の {time} に予約しました"


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