"""xAI Grok model configurations."""

from typing import Dict

from .base import ModelConfig, ParameterConfig, ParameterType

# xAI uses OpenAI-compatible API with max_tokens parameter

XAI_MODELS: Dict[str, ModelConfig] = {
    "grok-beta": ModelConfig(
        model_identifier="grok-beta",
        name="grok-beta",
        description="xAI's Grok beta model with strong reasoning and real-time knowledge.",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=True,
        supports_structured_outputs=False,
        reasoning=False,
        maximum_context_tokens=131072,
        maximum_output_tokens=4096,
        token_param_name="max_tokens",
        supports_temperature=True,
        input_cost_hint=5.0,
        output_cost_hint=15.0,
    ),
    "grok-2": ModelConfig(
        model_identifier="grok-2",
        name="grok-2",
        description="xAI's Grok 2 model with improved capabilities.",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=True,
        supports_structured_outputs=False,
        reasoning=False,
        maximum_context_tokens=131072,
        maximum_output_tokens=4096,
        token_param_name="max_tokens",
        supports_temperature=True,
        input_cost_hint=2.0,
        output_cost_hint=10.0,
    ),
    "grok-2-mini": ModelConfig(
        model_identifier="grok-2-mini",
        name="grok-2-mini",
        description="xAI's smaller Grok 2 model - faster and more cost-effective.",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=True,
        supports_structured_outputs=False,
        reasoning=False,
        maximum_context_tokens=131072,
        maximum_output_tokens=4096,
        token_param_name="max_tokens",
        supports_temperature=True,
        input_cost_hint=0.40,
        output_cost_hint=2.0,
    ),
    "grok-vision-beta": ModelConfig(
        model_identifier="grok-vision-beta",
        name="grok-vision-beta",
        description="xAI's Grok model with vision capabilities.",
        support_images=True,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=True,
        supports_structured_outputs=False,
        reasoning=False,
        maximum_context_tokens=8192,
        maximum_output_tokens=4096,
        token_param_name="max_tokens",
        supports_temperature=True,
        input_cost_hint=5.0,
        output_cost_hint=15.0,
    ),
}


XAI_PARAMETERS: list[ParameterConfig] = [
    ParameterConfig(
        field_name="temperature",
        display_name="Temperature",
        description="Controls randomness in responses. Higher values make output more random.",
        parameter_type=ParameterType.NUMBER,
        default_value=0.7,
        min_value=0,
        max_value=2,
        step=0.1,
    ),
    ParameterConfig(
        field_name="top_p",
        display_name="Top P",
        description="Nucleus sampling parameter.",
        parameter_type=ParameterType.NUMBER,
        default_value=1.0,
        min_value=0,
        max_value=1,
        step=0.1,
    ),
    ParameterConfig(
        field_name="max_tokens",
        display_name="Max Tokens",
        description="Maximum number of tokens to generate.",
        parameter_type=ParameterType.NUMBER,
        default_value=4096,
        min_value=1,
        max_value=4096,
        step=4,
    ),
    ParameterConfig(
        field_name="frequency_penalty",
        display_name="Frequency Penalty",
        description="Penalizes tokens based on their frequency in the text.",
        parameter_type=ParameterType.NUMBER,
        default_value=0,
        min_value=-2,
        max_value=2,
        step=0.1,
    ),
    ParameterConfig(
        field_name="presence_penalty",
        display_name="Presence Penalty",
        description="Penalizes tokens based on their presence in the text.",
        parameter_type=ParameterType.NUMBER,
        default_value=0,
        min_value=-2,
        max_value=2,
        step=0.1,
    ),
]
