from enum import Enum, StrEnum

from pydantic import BaseModel, Field

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

mcp = FastMCP(name="Elicitation Example")


class FruitEnum(str, Enum):
    pear = "pear"
    apple = "apple"


class TestStrEnum(StrEnum):
    one = "1"
    two = "2"
    three = "3"


class BookingPreferences(BaseModel):
    """Schema for collecting user preferences."""

    checkAlternative: bool = Field(description="Would you like to check another date?")
    alternativeDate: str = Field(
        default="2024-12-26",
        description="Alternative date (YYYY-MM-DD)",
    )
    fruit: FruitEnum = FruitEnum.pear
    strenum: TestStrEnum = TestStrEnum.one


@mcp.tool()
async def book_table(date: str, time: str, party_size: int, ctx: Context[ServerSession, None]) -> str:
    """Book a table with date availability check."""
    # Check if date is available
    if date == "2024-12-25":
        # Date unavailable - ask user for alternative
        result = await ctx.elicit(
            message=(f"No tables available for {party_size} on {date}. Would you like to try another date?"),
            schema=BookingPreferences,
        )
        print(result)
        if result.action == "accept" and result.data:
            if result.data.checkAlternative:
                return f"[SUCCESS] Booked for {result.data.alternativeDate}"
            return "[CANCELLED] No booking made"
        return "[CANCELLED] Booking cancelled"

    # Date available
    return f"[SUCCESS] Booked for {date} at {time}"


@mcp.tool()
async def get_uuid(ctx: Context[ServerSession, None]) -> str:
    """Get a unique identifier."""
    result = await ctx.elicit(
        message="Please enter a unique identifier",
        schema=InvalidFormatSchema,
    )
    if result.action == "accept" and result.data:
        return f"[SUCCESS] Unique identifier: {result.data.id}"
    return "[CANCELLED] No unique identifier provided"


if __name__ == "__main__":
    mcp.run(transport="streamable-http")