"""Data classes for Yera's config."""

from typing import Literal

from pydantic import BaseModel, ConfigDict, Field

MODEL_TYPES = ["llm", "embedding", "reranker", "image", "moderation", "tts", "stt"]

# ============================================================================
# Shared Components
# ============================================================================


class ModelProperties(BaseModel):
    """Technical properties of a model."""

    remote: bool = True
    context_window: int | None = None  # LLMs only

    model_config = ConfigDict(extra="forbid")


# ============================================================================
# Global Defaults
# ============================================================================


class GlobalDefaults(BaseModel):
    """Global default settings for Yera."""

    max_retries: int = 3
    timeout_seconds: int = 60
    log_level: Literal["debug", "info", "warn", "error"] = "info"
    enable_telemetry: bool = False

    model_config = ConfigDict(extra="forbid")


# ============================================================================
# LLM Models
# ============================================================================


class LLMCapabilities(BaseModel):
    """Capabilities specific to LLM models."""

    function_calling: bool = False
    structured_generation: bool = False
    vision: bool = False
    reasoning: bool = False
    multilingual: bool = False

    model_config = ConfigDict(extra="forbid")


class LLMCost(BaseModel):
    """Cost structure for LLM models."""

    input_tokens_per_million: float
    output_tokens_per_million: float

    model_config = ConfigDict(extra="forbid")


class LLMInference(BaseModel):
    """Runtime inference parameters for LLM models."""

    temperature: float | None = None
    max_tokens: int | None = None
    top_p: float | None = None
    top_k: int | None = None
    stop_sequences: list[str] = Field(default_factory=list)

    model_config = ConfigDict(extra="allow")


class LLMConfig(BaseModel):
    """Configuration for a Language Model."""

    id: str  # Full identifier like "anthropic.claude-sonnet-4"
    display_name: str
    credentials: str
    provider: str
    interface: str
    model_id: str
    properties: ModelProperties = Field(default_factory=ModelProperties)
    cost: LLMCost | None = None
    capabilities: LLMCapabilities = Field(default_factory=LLMCapabilities)
    inference: LLMInference = Field(default_factory=LLMInference)

    model_config = ConfigDict(extra="forbid")


# ============================================================================
# Embedding Models
# ============================================================================


class EmbeddingCapabilities(BaseModel):
    """Capabilities specific to embedding models."""

    cosine_similarity: bool = True
    dot_product: bool = True
    euclidean_distance: bool = False

    model_config = ConfigDict(extra="forbid")


class EmbeddingCost(BaseModel):
    """Cost structure for embedding models."""

    tokens_per_million: float

    model_config = ConfigDict(extra="forbid")


class EmbeddingInference(BaseModel):
    """Runtime inference parameters for embedding models."""

    batch_size: int | None = None
    dimensions: int | None = None  # For models that support configurable dimensions

    model_config = ConfigDict(extra="forbid")


class EmbeddingModelConfig(BaseModel):
    """Configuration for an embedding model."""

    id: str  # Full identifier like "openai.text-embedding-3-small"
    credentials: str
    provider: str
    interface: str
    model_id: str
    dimensions: int
    max_input_tokens: int
    status: Literal["stable", "beta", "deprecated"] = "stable"
    properties: ModelProperties
    cost: EmbeddingCost
    capabilities: EmbeddingCapabilities = Field(default_factory=EmbeddingCapabilities)
    inference: EmbeddingInference = Field(default_factory=EmbeddingInference)

    model_config = ConfigDict(extra="forbid")


# ============================================================================
# Reranker Models
# ============================================================================


class RerankerCapabilities(BaseModel):
    """Capabilities specific to reranker models."""

    multilingual: bool = False
    context_aware: bool = True

    model_config = ConfigDict(extra="forbid")


class RerankerCost(BaseModel):
    """Cost structure for reranker models."""

    per_thousand_searches: float

    model_config = ConfigDict(extra="forbid")


class RerankerInference(BaseModel):
    """Runtime inference parameters for reranker models."""

    top_n: int | None = None
    return_documents: bool = True

    model_config = ConfigDict(extra="forbid")


class RerankerModelConfig(BaseModel):
    """Configuration for a reranker model."""

    id: str  # Full identifier like "cohere.rerank-english-v3"
    credentials: str
    provider: str
    interface: str
    model_id: str
    max_documents: int
    max_query_length: int
    status: Literal["stable", "beta", "deprecated"] = "stable"
    properties: ModelProperties
    cost: RerankerCost
    capabilities: RerankerCapabilities = Field(default_factory=RerankerCapabilities)
    inference: RerankerInference = Field(default_factory=RerankerInference)

    model_config = ConfigDict(extra="forbid")


# ============================================================================
# Image Generation Models
# ============================================================================


class ImageCapabilities(BaseModel):
    """Capabilities specific to image generation models."""

    supports_variations: bool = False
    supports_edits: bool = False
    supports_transparent_background: bool = False

    model_config = ConfigDict(extra="forbid")


class ImageCost(BaseModel):
    """Cost structure for image generation models."""

    per_image_standard: float | None = None
    per_image_hd: float | None = None
    per_image_1024: float | None = None  # DALL-E 2 pricing
    per_image_512: float | None = None
    per_image_256: float | None = None

    model_config = ConfigDict(extra="forbid")


class ImageInference(BaseModel):
    """Runtime inference parameters for image generation models."""

    size: str | None = None
    quality: str | None = None
    style: str | None = None
    n: int = 1  # Number of images to generate

    model_config = ConfigDict(extra="forbid")


class ImageModelConfig(BaseModel):
    """Configuration for an image generation model."""

    id: str  # Full identifier like "openai.dall-e-3"
    credentials: str
    provider: str
    interface: str
    model_id: str
    supported_sizes: list[str]
    supported_qualities: list[str]
    supported_styles: list[str] = Field(default_factory=list)
    status: Literal["stable", "beta", "deprecated"] = "stable"
    properties: ModelProperties
    cost: ImageCost
    capabilities: ImageCapabilities = Field(default_factory=ImageCapabilities)
    inference: ImageInference = Field(default_factory=ImageInference)

    model_config = ConfigDict(extra="forbid")


# ============================================================================
# Moderation Models
# ============================================================================


class ModerationCapabilities(BaseModel):
    """Capabilities specific to moderation models."""

    categories: list[str] = Field(default_factory=list)

    model_config = ConfigDict(extra="forbid")


class ModerationCost(BaseModel):
    """Cost structure for moderation models."""

    free: bool = True

    model_config = ConfigDict(extra="forbid")


class ModerationInference(BaseModel):
    """Runtime inference parameters for moderation models."""

    # Most moderation models have no configurable inference params
    model_config = ConfigDict(extra="forbid")


class ModerationModelConfig(BaseModel):
    """Configuration for a moderation model."""

    id: str  # Full identifier like "openai.text-moderation-latest"
    credentials: str
    provider: str
    interface: str
    model_id: str
    status: Literal["stable", "beta", "deprecated"] = "stable"
    properties: ModelProperties
    cost: ModerationCost
    capabilities: ModerationCapabilities = Field(default_factory=ModerationCapabilities)
    inference: ModerationInference = Field(default_factory=ModerationInference)

    model_config = ConfigDict(extra="forbid")


# ============================================================================
# Text-to-Speech Models
# ============================================================================


class TTSCapabilities(BaseModel):
    """Capabilities specific to TTS models."""

    streaming: bool = False
    speed_control: bool = False
    voice_cloning: bool = False
    emotion_control: bool = False

    model_config = ConfigDict(extra="forbid")


class TTSCost(BaseModel):
    """Cost structure for TTS models."""

    chars_per_million: float

    model_config = ConfigDict(extra="forbid")


class TTSInference(BaseModel):
    """Runtime inference parameters for TTS models."""

    voice: str | None = None
    speed: float | None = None
    format: str | None = None  # mp3, opus, etc.
    sample_rate: int | None = None

    model_config = ConfigDict(extra="forbid")


class TTSModelConfig(BaseModel):
    """Configuration for a text-to-speech model."""

    id: str  # Full identifier like "openai.tts-1"
    credentials: str
    provider: str
    interface: str
    model_id: str
    voices: list[str] = Field(default_factory=list)
    default_voice: str = ""
    supported_formats: list[str] = Field(default_factory=list)
    sample_rates: list[int] = Field(default_factory=list)
    status: Literal["stable", "beta", "deprecated"] = "stable"
    properties: ModelProperties
    cost: TTSCost
    capabilities: TTSCapabilities = Field(default_factory=TTSCapabilities)
    inference: TTSInference = Field(default_factory=TTSInference)

    model_config = ConfigDict(extra="forbid")


# ============================================================================
# Speech-to-Text Models
# ============================================================================


class STTCapabilities(BaseModel):
    """Capabilities specific to STT models."""

    timestamps: bool = False
    translation: bool = False
    language_detection: bool = False

    model_config = ConfigDict(extra="forbid")


class STTCost(BaseModel):
    """Cost structure for STT models."""

    per_minute: float

    model_config = ConfigDict(extra="forbid")


class STTInference(BaseModel):
    """Runtime inference parameters for STT models."""

    language: str | None = None
    enable_timestamps: bool = False
    response_format: str | None = None  # json, text, srt, vtt

    model_config = ConfigDict(extra="forbid")


class STTModelConfig(BaseModel):
    """Configuration for a speech-to-text model."""

    id: str  # Full identifier like "openai.whisper-1"
    credentials: str
    provider: str
    interface: str
    model_id: str
    supported_formats: list[str] = Field(default_factory=list)
    max_file_size_mb: float = 25.0
    languages: list[str] = Field(default_factory=list)
    status: Literal["stable", "beta", "deprecated"] = "stable"
    properties: ModelProperties
    cost: STTCost
    capabilities: STTCapabilities = Field(default_factory=STTCapabilities)
    inference: STTInference = Field(default_factory=STTInference)

    model_config = ConfigDict(extra="forbid")


# ============================================================================
# Model Registry
# ============================================================================


class ModelRegistry(BaseModel):
    """Registry of all configured models by type.

    Models are keyed by their full dot-path identifier (e.g.,
    "anthropic.claude-sonnet-4").

    """

    llm: dict[str, LLMConfig] = Field(default_factory=dict)
    embedding: dict[str, EmbeddingModelConfig] = Field(default_factory=dict)
    reranker: dict[str, RerankerModelConfig] = Field(default_factory=dict)
    image: dict[str, ImageModelConfig] = Field(default_factory=dict)
    moderation: dict[str, ModerationModelConfig] = Field(default_factory=dict)
    tts: dict[str, TTSModelConfig] = Field(default_factory=dict)
    stt: dict[str, STTModelConfig] = Field(default_factory=dict)

    model_config = ConfigDict(extra="forbid")


# ============================================================================
# Top-Level Configuration
# ============================================================================


class DefaultModels(BaseModel):
    """Default model selections by type."""

    llm: str | None = None
    embedding: str | None = None
    reranker: str | None = None
    image: str | None = None
    tts: str | None = None
    stt: str | None = None
    moderation: str | None = None

    model_config = ConfigDict(extra="forbid")


class Defaults(BaseModel):
    """Yera configuration defaults section."""

    max_retries: int = 3
    timeout_seconds: int = 60
    log_level: Literal["debug", "info", "warn", "error"] = "info"
    enable_telemetry: bool = False

    models: DefaultModels = Field(default_factory=DefaultModels)

    model_config = ConfigDict(extra="forbid")


class CredentialsMap(BaseModel):
    """Credentials for Yera configuration."""

    providers: dict[str, dict[str, str]] = Field(default_factory=dict)
    tools: dict[str, dict[str, str]] = Field(default_factory=dict)


class YeraConfig(BaseModel):
    """Complete Yera configuration.

    This is the root configuration object that contains all provider credentials,
    model definitions, and default selections.
    """

    defaults: Defaults = Field(default_factory=Defaults)
    credentials: CredentialsMap = Field(default_factory=CredentialsMap)
    models: ModelRegistry = Field(default_factory=ModelRegistry)

    model_config = ConfigDict(extra="forbid")

    def merge(self, other: "YeraConfig") -> "YeraConfig":
        """Merge this configuration with another, creating a new config.

        Creates a new YeraConfig by combining this configuration with another,
        where values from `other` take precedence in case of conflicts. This is
        for extending and updating yera config.

        Args:
            other: Another YeraConfig to merge with this one. Values from this
                config will override values from self when keys conflict.

        Returns:
            A new YeraConfig instance containing the merged configuration. The
            original configs (self and other) are not modified.

        """
        merged_defaults = {
            **self.defaults.model_dump(),
            **other.defaults.model_dump(),
        }
        merged_default_models = {
            **self.defaults.models.model_dump(),
            **other.defaults.models.model_dump(),
        }
        defaults = self.defaults.model_validate(
            {
                **merged_defaults,
                "models": merged_default_models,
            }
        )

        providers_merge = {**self.credentials.providers, **other.credentials.providers}
        tools_merge = {**self.credentials.tools, **other.credentials.tools}
        credentials = self.credentials.model_copy(
            update={
                "providers": providers_merge,
                "tools": tools_merge,
            }
        )

        model_reg1 = self.models.model_dump()
        model_reg2 = other.models.model_dump()
        models_merge = {}
        for k in model_reg1:
            models_merge[k] = {**model_reg1[k], **model_reg2[k]}
        models = ModelRegistry(**models_merge)

        return YeraConfig(
            defaults=defaults,
            credentials=credentials,
            models=models,
        )
