from google.genai.types import ThinkingConfigDict, ThinkingLevel
from pydantic_ai.models.google import GoogleModel, GoogleModelSettings
from pydantic_ai.providers.google import GoogleProvider


def _thinking_config(model_name: str) -> ThinkingConfigDict:
    if model_name.startswith(("gemini-3.1-flash", "gemini-3-flash")):
        return ThinkingConfigDict(
            include_thoughts=False,
            thinking_level=ThinkingLevel.MINIMAL,
        )

    elif model_name.startswith(("gemini-3.1", "gemini-3")):
        return ThinkingConfigDict(
            include_thoughts=False,
            thinking_level=ThinkingLevel.LOW,
        )

    else:
        return ThinkingConfigDict(
            include_thoughts=False,
            thinking_budget=0,
        )


def build_google_model(model_name: str) -> GoogleModel:
    provider = GoogleProvider()
    thinking_config = _thinking_config(model_name)
    model_settings = GoogleModelSettings(google_thinking_config=thinking_config)
    return GoogleModel(model_name=model_name, provider=provider, settings=model_settings)
