"""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-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 4.1 Fast models (latest flagships)
    "grok-4-1-fast-reasoning": ModelConfig(
        model_identifier="grok-4-1-fast-reasoning",
        name="grok-4-1-fast-reasoning",
        description="xAI's Grok 4.1 Fast with reasoning - 2M context, agentic tool calling.",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=True,
        supports_structured_outputs=False,
        reasoning=True,
        maximum_context_tokens=2000000,
        maximum_output_tokens=16384,
        token_param_name="max_tokens",
        supports_temperature=True,
        input_cost_hint=0.20,
        output_cost_hint=0.50,
    ),
    "grok-4-1-fast-non-reasoning": ModelConfig(
        model_identifier="grok-4-1-fast-non-reasoning",
        name="grok-4-1-fast-non-reasoning",
        description="xAI's Grok 4.1 Fast without reasoning - low-latency instant responses.",
        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=2000000,
        maximum_output_tokens=16384,
        token_param_name="max_tokens",
        supports_temperature=True,
        input_cost_hint=0.20,
        output_cost_hint=0.50,
    ),
    # Grok Code model
    "grok-code-fast-1": ModelConfig(
        model_identifier="grok-code-fast-1",
        name="grok-code-fast-1",
        description="xAI's Grok optimized for agentic coding tasks.",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=True,
        supports_structured_outputs=False,
        reasoning=True,
        maximum_context_tokens=262144,
        maximum_output_tokens=16384,
        token_param_name="max_tokens",
        supports_temperature=True,
        input_cost_hint=0.10,
        output_cost_hint=0.30,
    ),
    # Grok 4 Fast models
    "grok-4-fast-reasoning": ModelConfig(
        model_identifier="grok-4-fast-reasoning",
        name="grok-4-fast-reasoning",
        description="xAI's Grok 4 Fast with reasoning capabilities.",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=True,
        supports_structured_outputs=False,
        reasoning=True,
        maximum_context_tokens=131072,
        maximum_output_tokens=16384,
        token_param_name="max_tokens",
        supports_temperature=True,
        input_cost_hint=0.50,
        output_cost_hint=1.50,
    ),
    "grok-4-fast-non-reasoning": ModelConfig(
        model_identifier="grok-4-fast-non-reasoning",
        name="grok-4-fast-non-reasoning",
        description="xAI's Grok 4 Fast without reasoning - low-latency responses.",
        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=16384,
        token_param_name="max_tokens",
        supports_temperature=True,
        input_cost_hint=0.50,
        output_cost_hint=1.50,
    ),
    "grok-4-0709": ModelConfig(
        model_identifier="grok-4-0709",
        name="grok-4-0709",
        description="xAI's Grok 4 model (July 2025 version).",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=True,
        supports_structured_outputs=False,
        reasoning=True,
        maximum_context_tokens=131072,
        maximum_output_tokens=16384,
        token_param_name="max_tokens",
        supports_temperature=True,
        input_cost_hint=2.0,
        output_cost_hint=6.0,
    ),
    # Grok 3 models
    "grok-3": ModelConfig(
        model_identifier="grok-3",
        name="grok-3",
        description="xAI's Grok 3 model with advanced reasoning.",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=True,
        supports_structured_outputs=False,
        reasoning=True,
        maximum_context_tokens=131072,
        maximum_output_tokens=8192,
        token_param_name="max_tokens",
        supports_temperature=True,
        input_cost_hint=3.0,
        output_cost_hint=15.0,
    ),
    "grok-3-mini": ModelConfig(
        model_identifier="grok-3-mini",
        name="grok-3-mini",
        description="xAI's Grok 3 mini - faster and cost-effective.",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=True,
        supports_structured_outputs=False,
        reasoning=True,
        maximum_context_tokens=131072,
        maximum_output_tokens=8192,
        token_param_name="max_tokens",
        supports_temperature=True,
        input_cost_hint=0.30,
        output_cost_hint=0.50,
    ),
    # Grok 2 Vision
    "grok-2-vision-1212": ModelConfig(
        model_identifier="grok-2-vision-1212",
        name="grok-2-vision-1212",
        description="xAI's Grok 2 Vision model with image support.",
        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=32768,
        maximum_output_tokens=4096,
        token_param_name="max_tokens",
        supports_temperature=True,
        input_cost_hint=2.0,
        output_cost_hint=10.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,
    ),
]
