"""Data models for AI models configuration."""

from __future__ import annotations

from datetime import datetime
from enum import Enum, StrEnum
from typing import Any

from pydantic import BaseModel, ConfigDict, Field

from ..core.enums import ModelProvider, ModelPublisher


def _empty_variant_list() -> list[BaseModelConfigVariant]:
    return []


def _empty_model_uses_provider_list() -> list[BaseModelUsesProvider]:
    return []


class ModelOutputFormat(StrEnum):
    """Enum for ModelOutputFormat."""

    TOOL_CALL_JSON = "ModelOutputFormatToolCallJson"
    XML = "ModelOutputFormatXml"


class ReasoningEffort(StrEnum):
    """Enum for ReasoningEffort."""

    LOW = "ReasoningEffortLow"
    MEDIUM = "ReasoningEffortMedium"
    HIGH = "ReasoningEffortHigh"


class ModelCompatibility(BaseModel):
    """Data contract for ModelCompatibility."""

    has_image_support: bool = Field(alias="hasImageSupport")


class BaseModelShared(BaseModel):
    """Data contract for BaseModelShared."""

    default_max_convo_tokens: int = Field(alias="defaultMaxConvoTokens")
    max_tokens: int = Field(alias="maxTokens")
    max_output_tokens: int = Field(alias="maxOutputTokens")
    reserved_output_tokens: int = Field(alias="reservedOutputTokens")
    preferred_output_format: ModelOutputFormat = Field(alias="preferredOutputFormat")
    system_prompt_disabled: bool | None = Field(
        default=None, alias="systemPromptDisabled"
    )
    role_params_disabled: bool | None = Field(default=None, alias="roleParamsDisabled")
    stop_disabled: bool | None = Field(default=None, alias="stopDisabled")
    predicted_output_enabled: bool | None = Field(
        default=None, alias="predictedOutputEnabled"
    )
    reasoning_effort_enabled: bool | None = Field(
        default=None, alias="reasoningEffortEnabled"
    )
    reasoning_effort: ReasoningEffort | None = Field(
        default=None, alias="reasoningEffort"
    )
    include_reasoning: bool | None = Field(default=None, alias="includeReasoning")
    hide_reasoning: bool | None = Field(default=None, alias="hideReasoning")
    reasoning_budget: int | None = Field(default=None, alias="reasoningBudget")
    supports_cache_control: bool | None = Field(
        default=None, alias="supportsCacheControl"
    )
    single_message_no_system_prompt: bool | None = Field(
        default=None, alias="singleMessageNoSystemPrompt"
    )
    token_estimate_padding_pct: float | None = Field(
        default=None, alias="tokenEstimatePaddingPct"
    )


class BaseModelProviderConfig(BaseModel):
    """Data contract for BaseModelProviderConfig."""

    model_name: str = Field(alias="modelName")


class BaseModelConfig(BaseModel):
    """Base model configuration."""

    model_tag: str = Field(alias="modelTag")
    model_id: str = Field(alias="modelId")
    publisher: ModelPublisher | None = Field(alias="publisher", default=None)
    base_model_shared: BaseModelProviderConfig = Field(alias="basemodelshared")

    model_config = ConfigDict(
        str_strip_whitespace=True,
        validate_assignment=True,
        arbitrary_types_allowed=False,
        populate_by_name=True,
        use_enum_values=True,
    )


class BaseModelUsesProvider(BaseModel):
    """Data contract for BaseModelUsesProvider."""

    provider: ModelProvider
    custom_provider: str | None = Field(default=None, alias="customProvider")
    model_name: str = Field(alias="modelName")


class BaseModelConfigVariant(BaseModel):
    """Data contract for BaseModelConfigVariant."""

    is_base_variant: bool = Field(alias="isBaseVariant")
    variant_tag: str = Field(alias="variantTag")
    description: str
    overrides: BaseModelShared
    variants: list[BaseModelConfigVariant] = Field(default_factory=_empty_variant_list)
    requires_variant_overrides: list[str] = Field(
        default_factory=list, alias="requiresVariantOverrides"
    )
    is_default_variant: bool = Field(alias="isDefaultVariant")


class BaseModelConfigSchema(BaseModel):
    """Data contract for BaseModelConfigSchema."""

    model_tag: str = Field(alias="modelTag")
    model_id: str = Field(alias="modelId")
    publisher: ModelPublisher
    description: str
    requires_variant_overrides: list[str] = Field(
        default_factory=list, alias="requiresVariantOverrides"
    )
    variants: list[BaseModelConfigVariant] = Field(default_factory=_empty_variant_list)
    providers: list[BaseModelUsesProvider] = Field(
        default_factory=_empty_model_uses_provider_list
    )


class AvailableModel(BaseModel):
    """Data contract for AvailableModel."""

    id: str
    description: str
    default_max_convo_tokens: int = Field(alias="defaultMaxConvoTokens")
    created_at: datetime = Field(alias="createdAt")
    updated_at: datetime = Field(alias="updatedAt")


class PlannerModelConfig(BaseModel):
    """Data contract for PlannerModelConfig."""

    max_convo_tokens: int = Field(alias="maxConvoTokens")


class ModelRoleConfig(BaseModel):
    """Model role configuration."""

    role: str
    model_id: str = Field(alias="modelId")
    base_model_config: BaseModelConfig | None = Field(
        alias="baseModelConfig", default=None
    )
    temperature: float
    top_p: float = Field(alias="topP")
    reserved_output_tokens: int = Field(alias="reservedOutputTokens")
    large_context_fallback: ModelRoleConfig | None = Field(
        default=None, alias="largeContextFallback"
    )
    large_output_fallback: ModelRoleConfig | None = Field(
        default=None, alias="largeOutputFallback"
    )
    error_fallback: ModelRoleConfig | None = Field(default=None, alias="errorFallback")
    strong_model: ModelRoleConfig | None = Field(default=None, alias="strongModel")
    local_provider: ModelProvider | None = Field(default=None, alias="localProvider")

    model_config = ConfigDict(
        str_strip_whitespace=True,
        validate_assignment=True,
        arbitrary_types_allowed=False,
        populate_by_name=True,
        use_enum_values=True,
    )


class ModelRoleModelConfig(BaseModel):
    """Data contract for ModelRoleModelConfig."""

    provider: ModelProvider
    custom_provider: str | None = Field(default=None, alias="customProvider")
    model_tag: str = Field(alias="modelTag")


class ModelRoleConfigSchema(BaseModel):
    """Data contract for ModelRoleConfigSchema."""

    model_id: str = Field(alias="modelId")
    temperature: float | None = None
    top_p: float | None = Field(default=None, alias="topP")
    reserved_output_tokens: int | None = Field(
        default=None, alias="reservedOutputTokens"
    )
    max_convo_tokens: int | None = Field(default=None, alias="maxConvoTokens")
    large_context_fallback: ModelRoleConfigSchema | None = Field(
        default=None, alias="largeContextFallback"
    )
    large_output_fallback: ModelRoleConfigSchema | None = Field(
        default=None, alias="largeOutputFallback"
    )
    error_fallback: ModelRoleConfigSchema | None = Field(
        default=None, alias="errorFallback"
    )
    strong_model: ModelRoleConfigSchema | None = Field(
        default=None, alias="strongModel"
    )


class PlannerRoleConfig(BaseModel):
    """Data contract for PlannerRoleConfig."""

    model_role_config: PlannerModelConfig = Field(alias="modelroleconfig")


class ClientModelPackSchemaRoles(BaseModel):
    """Data contract for ClientModelPackSchemaRoles."""

    schema_url: Any | None = Field(default=None, alias="$schema")
    local_provider: ModelProvider | None = Field(default=None, alias="localProvider")
    planner: Any
    architect: Any | None = None
    coder: Any | None = None
    summarizer: Any
    builder: Any
    whole_file_builder: Any | None = Field(default=None, alias="wholeFileBuilder")
    names: Any
    commit_messages: Any = Field(alias="commitMessages")
    auto_continue: Any = Field(alias="autoContinue")


class ModelPackSchemaRoles(BaseModel):
    """Data contract for ModelPackSchemaRoles."""

    local_provider: ModelProvider | None = Field(default=None, alias="localProvider")
    planner: ModelRoleConfigSchema
    coder: ModelRoleConfigSchema | None = None
    plan_summary: ModelRoleConfigSchema = Field(alias="planSummary")
    builder: ModelRoleConfigSchema
    whole_file_builder: ModelRoleConfigSchema | None = Field(
        default=None, alias="wholeFileBuilder"
    )
    namer: ModelRoleConfigSchema
    commit_msg: ModelRoleConfigSchema = Field(alias="commitMsg")
    exec_status: ModelRoleConfigSchema = Field(alias="execStatus")
    context_loader: ModelRoleConfigSchema | None = Field(
        default=None, alias="contextLoader"
    )


class ModelPackSchema(BaseModel):
    """Data contract for ModelPackSchema."""

    name: str
    description: str


class ModelPack(BaseModel):
    """Data contract for ModelPack."""

    id: str
    name: str
    local_provider: ModelProvider | None = Field(default=None, alias="localProvider")
    description: str
    planner: PlannerRoleConfig
    coder: ModelRoleConfig | None = None
    plan_summary: ModelRoleConfig = Field(alias="planSummary")
    builder: ModelRoleConfig
    whole_file_builder: ModelRoleConfig | None = Field(
        default=None, alias="wholeFileBuilder"
    )
    namer: ModelRoleConfig
    commit_msg: ModelRoleConfig = Field(alias="commitMsg")
    exec_status: ModelRoleConfig = Field(alias="execStatus")
    context_loader: ModelRoleConfig | None = Field(default=None, alias="contextLoader")


class AIModel(BaseModel):
    """Data contract for AIModel."""

    provider_id: str = Field(alias="providerId")
    model_id: str = Field(alias="modelId")
    is_chat: bool = Field(default=False, alias="isChat")
    is_image_to_text: bool = Field(default=False, alias="isImageToText")
    is_multi_modal: bool = Field(default=False, alias="isMultiModal")
    is_fine_tunable: bool = Field(default=False, alias="isFineTunable")
    is_function_calling: bool = Field(default=False, alias="isFunctionCalling")

    model_config = ConfigDict(populate_by_name=True)


class Provider(BaseModel):
    """Data contract for Provider."""

    provider_id: str = Field(alias="providerId")
    name: str

    model_config = ConfigDict(populate_by_name=True)
