import asyncio

from dotenv import load_dotenv
from fastmcp.client.sampling.handlers.openai import OpenAISamplingHandler
from mcp.types import CreateMessageRequestParams, SamplingMessage, TextContent

load_dotenv()


async def main():

    handler = OpenAISamplingHandler(default_model="gpt-4o")
    messages = [
        SamplingMessage(role="user", content=TextContent(type="text", text="Hi"))
    ]
    params = CreateMessageRequestParams(messages=messages, maxTokens=300)
    result = await handler(messages, params, context=None)  # type: ignore
    print("LLM Response:", result)

    # ===> Hello! How can I assist you today?

    models = [
        "o1",
        "gpt-5-mini",
        "gpt-5",
    ]
    for model in models:
        try:
            handler = OpenAISamplingHandler(default_model=model)
            messages = [
                SamplingMessage(
                    role="user", content=TextContent(type="text", text="Hi")
                )
            ]
            params = CreateMessageRequestParams(messages=messages, maxTokens=300)
            result = await handler(messages, params, context=None)  # type: ignore
            print("LLM Response:", result)
        except Exception as e:
            print(f"Error with model {model}: {e}")

    # ===> 'max_tokens' is not supported with this model.
    # Use 'max_completion_tokens' instead.


# The problem:
# server/.venv/lib/python3.12/site-packages/fastmcp/client/sampling/handlers/openai.py
# blindly uses max_tokens for ALL models


asyncio.run(main())