from typing import Literal, Optional

from pydantic import BaseModel, Field


class EmbeddingConfig(BaseModel):
    """

    Embedding model configuration. This object specifies all the information necessary to access an embedding model to usage with Mirix, except for secret keys.

    Attributes:
        embedding_endpoint_type (str): The endpoint type for the model.
        embedding_endpoint (str): The endpoint for the model.
        embedding_model (str): The model for the embedding.
        embedding_dim (int): The dimension of the embedding.
        embedding_chunk_size (int): The chunk size of the embedding.
        auth_provider (str, optional): Name of registered auth provider for dynamic header injection.
        azure_endpoint (:obj:`str`, optional): The Azure endpoint for the model (Azure only).
        azure_version (str): The Azure version for the model (Azure only).
        azure_deployment (str): The Azure deployment for the model (Azure only).

    """

    embedding_endpoint_type: Literal[
        "openai",
        "anthropic",
        "bedrock",
        "cohere",
        "google_ai",
        "azure",
        "groq",
        "ollama",
        "webui",
        "webui-legacy",
        "lmstudio",
        "lmstudio-legacy",
        "llamacpp",
        "koboldcpp",
        "vllm",
        "hugging-face",
        "mistral",
        "together",  # completions endpoint
    ] = Field(..., description="The endpoint type for the model.")
    embedding_endpoint: Optional[str] = Field(
        None, description="The endpoint for the model (`None` if local)."
    )
    embedding_model: str = Field(..., description="The model for the embedding.")
    embedding_dim: int = Field(..., description="The dimension of the embedding.")
    embedding_chunk_size: Optional[int] = Field(
        300, description="The chunk size of the embedding."
    )
    handle: Optional[str] = Field(
        None,
        description="The handle for this config, in the format provider/model-name.",
    )
    auth_provider: Optional[str] = Field(
        None,
        description="Name of registered auth provider for dynamic header injection (e.g., for claims-based tickets)",
    )

    # azure only
    azure_endpoint: Optional[str] = Field(
        None, description="The Azure endpoint for the model."
    )
    azure_version: Optional[str] = Field(
        None, description="The Azure version for the model."
    )
    azure_deployment: Optional[str] = Field(
        None, description="The Azure deployment for the model."
    )

    @classmethod
    def default_config(
        cls, model_name: Optional[str] = None, provider: Optional[str] = None
    ):
        if model_name == "text-embedding-3-small" or (
            not model_name and provider == "openai"
        ):
            return cls(
                embedding_model="text-embedding-3-small",
                embedding_endpoint_type="openai",
                embedding_endpoint="https://api.openai.com/v1",
                embedding_dim=1536,
                embedding_chunk_size=8191,
            )
        elif model_name == "text-embedding-004" or (
            not model_name and provider == "google_ai"
        ):
            return cls(
                embedding_model="text-embedding-004",
                embedding_endpoint_type="google_ai",
                embedding_endpoint="https://generativelanguage.googleapis.com",
                embedding_dim=768,
                embedding_chunk_size=2048,
            )
        else:
            raise ValueError(f"Model {model_name} not supported.")

    def pretty_print(self) -> str:
        return (
            f"{self.embedding_model}"
            + (
                f" [type={self.embedding_endpoint_type}]"
                if self.embedding_endpoint_type
                else ""
            )
            + (f" [ip={self.embedding_endpoint}]" if self.embedding_endpoint else "")
        )
