"""Ollama model configurations for local models."""

from typing import Dict

from .base import ModelConfig, ParameterConfig, ParameterType

# Ollama uses num_predict parameter for max tokens
# Models need to be pulled locally first

OLLAMA_MODELS: Dict[str, ModelConfig] = {
    # Llama models
    "llama3.1": ModelConfig(
        model_identifier="llama3.1",
        name="llama3.1",
        description="Meta's Llama 3.1 model for general-purpose tasks.",
        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="num_predict",
        supports_temperature=True,
        input_cost_hint=0.0,  # Local model - no cost
        output_cost_hint=0.0,
    ),
    "llama3.1:8b": ModelConfig(
        model_identifier="llama3.1:8b",
        name="llama3.1-8b",
        description="Llama 3.1 8B parameter model - efficient and fast.",
        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="num_predict",
        supports_temperature=True,
        input_cost_hint=0.0,
        output_cost_hint=0.0,
    ),
    "llama3.1:70b": ModelConfig(
        model_identifier="llama3.1:70b",
        name="llama3.1-70b",
        description="Llama 3.1 70B parameter model - high performance.",
        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="num_predict",
        supports_temperature=True,
        input_cost_hint=0.0,
        output_cost_hint=0.0,
    ),
    "llama3.2": ModelConfig(
        model_identifier="llama3.2",
        name="llama3.2",
        description="Latest Llama 3.2 model with improved 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=131072,
        maximum_output_tokens=4096,
        token_param_name="num_predict",
        supports_temperature=True,
        input_cost_hint=0.0,
        output_cost_hint=0.0,
    ),
    # Mistral models
    "mistral": ModelConfig(
        model_identifier="mistral",
        name="mistral",
        description="Mistral 7B model - efficient and capable.",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=False,
        supports_structured_outputs=False,
        reasoning=False,
        maximum_context_tokens=32768,
        maximum_output_tokens=4096,
        token_param_name="num_predict",
        supports_temperature=True,
        input_cost_hint=0.0,
        output_cost_hint=0.0,
    ),
    "mixtral": ModelConfig(
        model_identifier="mixtral",
        name="mixtral",
        description="Mixtral 8x7B MoE model - excellent performance.",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=False,
        supports_structured_outputs=False,
        reasoning=False,
        maximum_context_tokens=32768,
        maximum_output_tokens=4096,
        token_param_name="num_predict",
        supports_temperature=True,
        input_cost_hint=0.0,
        output_cost_hint=0.0,
    ),
    # Code models
    "codellama": ModelConfig(
        model_identifier="codellama",
        name="codellama",
        description="Code Llama model specialized for coding tasks.",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=False,
        supports_structured_outputs=False,
        reasoning=False,
        maximum_context_tokens=16384,
        maximum_output_tokens=4096,
        token_param_name="num_predict",
        supports_temperature=True,
        input_cost_hint=0.0,
        output_cost_hint=0.0,
    ),
    # Other popular models
    "phi3": ModelConfig(
        model_identifier="phi3",
        name="phi3",
        description="Microsoft Phi-3 model - small but capable.",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=False,
        supports_structured_outputs=False,
        reasoning=False,
        maximum_context_tokens=4096,
        maximum_output_tokens=2048,
        token_param_name="num_predict",
        supports_temperature=True,
        input_cost_hint=0.0,
        output_cost_hint=0.0,
    ),
    "gemma2": ModelConfig(
        model_identifier="gemma2",
        name="gemma2",
        description="Google Gemma 2 open model.",
        support_images=False,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=False,
        supports_structured_outputs=False,
        reasoning=False,
        maximum_context_tokens=8192,
        maximum_output_tokens=4096,
        token_param_name="num_predict",
        supports_temperature=True,
        input_cost_hint=0.0,
        output_cost_hint=0.0,
    ),
    "qwen2.5": ModelConfig(
        model_identifier="qwen2.5",
        name="qwen2.5",
        description="Alibaba Qwen 2.5 model with multilingual support.",
        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=32768,
        maximum_output_tokens=4096,
        token_param_name="num_predict",
        supports_temperature=True,
        input_cost_hint=0.0,
        output_cost_hint=0.0,
    ),
    # Vision models
    "llava": ModelConfig(
        model_identifier="llava",
        name="llava",
        description="LLaVA vision-language model for image understanding.",
        support_images=True,
        support_files=False,
        support_streaming=True,
        supports_json_mode=True,
        supports_tool_call=False,
        supports_structured_outputs=False,
        reasoning=False,
        maximum_context_tokens=4096,
        maximum_output_tokens=2048,
        token_param_name="num_predict",
        supports_temperature=True,
        input_cost_hint=0.0,
        output_cost_hint=0.0,
    ),
}


OLLAMA_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.8,
        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=0.9,
        min_value=0,
        max_value=1,
        step=0.1,
    ),
    ParameterConfig(
        field_name="top_k",
        display_name="Top K",
        description="Top-k sampling parameter.",
        parameter_type=ParameterType.NUMBER,
        default_value=40,
        min_value=1,
        max_value=100,
        step=1,
    ),
    ParameterConfig(
        field_name="num_predict",
        display_name="Max Tokens",
        description="Maximum number of tokens to generate.",
        parameter_type=ParameterType.NUMBER,
        default_value=2048,
        min_value=1,
        max_value=4096,
        step=4,
    ),
    ParameterConfig(
        field_name="repeat_penalty",
        display_name="Repeat Penalty",
        description="Penalizes repetition in generated text.",
        parameter_type=ParameterType.NUMBER,
        default_value=1.1,
        min_value=0.1,
        max_value=2.0,
        step=0.1,
    ),
]


def get_token_param_name(model: str) -> str:
    """Get the correct token parameter name for an Ollama model.

    All Ollama models use num_predict.

    Args:
        model: The model identifier

    Returns:
        The API parameter name to use for max tokens
    """
    return "num_predict"
