# -------------------------------------------------------------------------------- #
# Model Capabilities
# -------------------------------------------------------------------------------- #
# This file is auto-generated by scripts/generate_model_capabilities.py
# DO NOT EDIT MANUALLY
# -------------------------------------------------------------------------------- #

# -------------------------------------------------------------------------------- #
# Imports
# -------------------------------------------------------------------------------- #
# Built-in imports
from typing import Literal, Dict, TypedDict, Union, Optional, NotRequired, Required

# Module imports
from astral_ai.constants._models import ModelAlias, ModelId, CompletionModels


# -------------------------------------------------------------------------------- #
# Feature Names
# -------------------------------------------------------------------------------- #

FEATURE_NAME = Literal[
    "developer_message",
    "distillation",
    "extended_thinking",
    "fine_tuning",
    "function_calling",
    "json_mode",
    "predicted_outputs",
    "reasoning_effort",
    "streaming",
    "structured_output",
    "system_message",
]


# -------------------------------------------------------------------------------- #
# Model Capabilities TypedDict
# -------------------------------------------------------------------------------- #

class ModelCapabilities(TypedDict, total=False):
    """TypedDict for capturing model capabilities."""
    supports_developer_message: Required[bool]
    supports_distillation: Required[bool]
    supports_extended_thinking: Required[bool]
    supports_fine_tuning: Required[bool]
    supports_function_calling: Required[bool]
    supports_json_mode: Required[bool]
    supports_predicted_outputs: Required[bool]
    supports_reasoning_effort: Required[bool]
    supports_streaming: Required[bool]
    supports_structured_output: Required[bool]
    supports_system_message: Required[bool]

    # Optional Capabilities 
    max_output_tokens: NotRequired[int]
    max_output_tokens_reasoning_effort: NotRequired[int | None]



# -------------------------------------------------------------------------------- #
# Alias to Model ID Mapping
# -------------------------------------------------------------------------------- #

ALIAS_TO_MODEL_ID: Dict[ModelAlias, ModelId] = {
    "claude-3-5-haiku": "claude-3-5-haiku-20241022",
    "claude-3-5-sonnet": "claude-3-5-sonnet-20241022",
    "claude-3-7-sonnet": "claude-3-7-sonnet-20250219",
    "claude-3-haiku": "claude-3-haiku-20240307",
    "claude-3-opus": "claude-3-opus-20240229",
    "claude-3-sonnet": "claude-3-sonnet-20240229",
    "deepseek-chat": "deepseek-chat",
    "deepseek-reasoner": "deepseek-reasoner",
    "gpt-4.5-preview": "gpt-4.5-preview-2025-02-27",
    "gpt-4o": "gpt-4o-2024-08-06",
    "gpt-4o-mini": "gpt-4o-mini-2024-07-18",
    "o1": "o1-2024-12-17",
    "o1-mini": "o1-mini-2024-09-12",
    "o1-preview": "o1-preview-2024-09-12",
    "o1-pro": "o1-pro-2025-03-19",
    "o3-mini": "o3-mini-2025-01-31",
}


# -------------------------------------------------------------------------------- #
# Model Capabilities Mapping
# -------------------------------------------------------------------------------- #

MODEL_CAPABILITIES: Dict[ModelId, ModelCapabilities] = {
    "claude-3-5-haiku-20241022": {
        "supports_developer_message": True,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": False,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": True,
        "supports_streaming": True,
        "supports_structured_output": False,
        "supports_system_message": False,
        "max_output_tokens": 8192,
        "max_output_tokens_reasoning_effort": None,
    },
    "claude-3-5-sonnet-20240620": {
        "supports_developer_message": True,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": False,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": True,
        "supports_streaming": True,
        "supports_structured_output": False,
        "supports_system_message": False,
        "max_output_tokens": 8192,
        "max_output_tokens_reasoning_effort": None,
    },
    "claude-3-5-sonnet-20241022": {
        "supports_developer_message": True,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": False,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": True,
        "supports_streaming": True,
        "supports_structured_output": False,
        "supports_system_message": False,
        "max_output_tokens": 8192,
        "max_output_tokens_reasoning_effort": None,
    },
    "claude-3-7-sonnet-20250219": {
        "supports_developer_message": True,
        "supports_distillation": False,
        "supports_extended_thinking": True,
        "supports_fine_tuning": False,
        "supports_function_calling": False,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": True,
        "supports_streaming": True,
        "supports_structured_output": False,
        "supports_system_message": False,
        "max_output_tokens": 8192,
        "max_output_tokens_reasoning_effort": 8191,
    },
    "claude-3-haiku-20240307": {
        "supports_developer_message": True,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": False,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": True,
        "supports_streaming": True,
        "supports_structured_output": False,
        "supports_system_message": False,
        "max_output_tokens": 4096,
        "max_output_tokens_reasoning_effort": None,
    },
    "claude-3-opus-20240229": {
        "supports_developer_message": True,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": False,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": True,
        "supports_streaming": True,
        "supports_structured_output": False,
        "supports_system_message": False,
        "max_output_tokens": 4096,
        "max_output_tokens_reasoning_effort": None,
    },
    "deepseek-chat": {
        "supports_developer_message": False,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": True,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": False,
        "supports_streaming": True,
        "supports_structured_output": False,
        "supports_system_message": True,
    },
    "deepseek-reasoner": {
        "supports_developer_message": True,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": False,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": True,
        "supports_streaming": True,
        "supports_structured_output": False,
        "supports_system_message": False,
        "max_output_tokens": 8000,
        "max_output_tokens_reasoning_effort": None,
    },
    "gpt-4.5-preview-2025-02-27": {
        "supports_developer_message": False,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": True,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": False,
        "supports_streaming": True,
        "supports_structured_output": True,
        "supports_system_message": True,
    },
    "gpt-4o-2024-05-13": {
        "supports_developer_message": False,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": True,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": False,
        "supports_streaming": True,
        "supports_structured_output": False,
        "supports_system_message": True,
    },
    "gpt-4o-2024-08-06": {
        "supports_developer_message": False,
        "supports_distillation": True,
        "supports_extended_thinking": False,
        "supports_fine_tuning": True,
        "supports_function_calling": True,
        "supports_json_mode": True,
        "supports_predicted_outputs": True,
        "supports_reasoning_effort": False,
        "supports_streaming": True,
        "supports_structured_output": True,
        "supports_system_message": True,
    },
    "gpt-4o-2024-11-20": {
        "supports_developer_message": False,
        "supports_distillation": True,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": True,
        "supports_json_mode": True,
        "supports_predicted_outputs": True,
        "supports_reasoning_effort": False,
        "supports_streaming": True,
        "supports_structured_output": True,
        "supports_system_message": True,
    },
    "gpt-4o-mini-2024-07-18": {
        "supports_developer_message": False,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": True,
        "supports_function_calling": True,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": False,
        "supports_streaming": True,
        "supports_structured_output": True,
        "supports_system_message": True,
    },
    "o1-2024-12-17": {
        "supports_developer_message": True,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": True,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": True,
        "supports_streaming": True,
        "supports_structured_output": True,
        "supports_system_message": False,
    },
    "o1-mini-2024-09-12": {
        "supports_developer_message": True,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": False,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": True,
        "supports_streaming": True,
        "supports_structured_output": False,
        "supports_system_message": False,
    },
    "o1-preview-2024-09-12": {
        "supports_developer_message": True,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": True,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": True,
        "supports_streaming": True,
        "supports_structured_output": True,
        "supports_system_message": False,
    },
    "o1-pro-2025-03-19": {
        "supports_developer_message": True,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": True,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": True,
        "supports_streaming": False,
        "supports_structured_output": True,
        "supports_system_message": False,
    },
    "o3-mini-2025-01-31": {
        "supports_developer_message": True,
        "supports_distillation": False,
        "supports_extended_thinking": False,
        "supports_fine_tuning": False,
        "supports_function_calling": True,
        "supports_json_mode": True,
        "supports_predicted_outputs": False,
        "supports_reasoning_effort": True,
        "supports_streaming": True,
        "supports_structured_output": True,
        "supports_system_message": False,
    },
}


# -------------------------------------------------------------------------------- #
# Helper Functions
# -------------------------------------------------------------------------------- #


def get_specific_model_id(model: Union[ModelAlias, ModelId]) -> ModelId:
    """
    Convert a model alias to its specific model ID.
    If a specific model ID is provided, return it directly.

    Args:
        model: The model alias or ID to convert

    Returns:
        ModelId: The specific model ID
    """
    if model in ALIAS_TO_MODEL_ID:
        return ALIAS_TO_MODEL_ID[model]
    return model  # It's already a specific model ID


def supports_feature(model: Union[ModelAlias, ModelId], feature: FEATURE_NAME) -> bool:
    """
    Check if a model supports a specific feature with O(1) lookup.
    Works with both aliases and specific model IDs.

    Args:
        model: The model alias or ID to check
        feature: The feature to check for (e.g., 'reasoning_effort')

    Returns:
        bool: True if the model supports the feature, False otherwise
    """
    model_id = get_specific_model_id(model)
    if model_id not in MODEL_CAPABILITIES:
        return False
    
    # Prepend 'supports_' to the feature name for lookup
    capability_key = f"supports_{feature}"
    return MODEL_CAPABILITIES[model_id].get(capability_key, False)




def get_model_max_tokens(model: Union[ModelAlias, ModelId], with_reasoning_effort: bool = False) -> Union[int, None]:
    """
    Get the maximum output tokens for a model.
    """
    model_id = get_specific_model_id(model)

    if model_id not in MODEL_CAPABILITIES:
        return None
    
    if MODEL_CAPABILITIES[model_id].get("supports_reasoning_effort", False) and with_reasoning_effort:
        return MODEL_CAPABILITIES[model_id].get("max_output_tokens_reasoning_effort", None)
    
    return MODEL_CAPABILITIES[model_id].get("max_output_tokens", None)
