from abc import ABC, abstractmethod
from typing import Optional, Any, Literal

from pydantic import BaseModel
from sqlmodel import Field


EmbeddingSourceType = Literal["text", "image"]


class BaseEmbeddingFunction(BaseModel, ABC):
    provider: str = Field("openai", description="The name of the embedding provider")
    model_name: str = Field(
        None, description="The name of embedding model used for embedding"
    )
    dimensions: Optional[int] = Field(
        None, description="Dimensions of the vector generated by the embedding model"
    )
    use_server: bool = Field(
        False,
        description=(
            "Whether to enable server-side embedding (auto embedding in the database). "
            "If True, the embedding for the source data will be computed in the database. "
            "Default is False."
        ),
    )
    server_embed_params: Optional[dict[str, Any]] = Field(
        None, description="The parameters for server-side embedding. "
    )

    def __init__(self, /, **data: Any):
        super().__init__(**data)

    def VectorField(
        self,
        source_field: Optional[str] = None,
        source_type: Optional[EmbeddingSourceType] = "text",
        **kwargs,
    ):
        """
        Create a VectorField with auto embedding configuration.

        Args:
            source_field: The name of the source field to embed
            source_type: The type of source data ("text" or "image")
            **kwargs: Additional keyword arguments for VectorField

        Returns:
            VectorField configured with this embedding function
        """
        # Import here to avoid circular import
        from pytidb.schema import VectorField

        return VectorField(
            embed_fn=self,
            dimensions=self.dimensions,
            source_field=source_field,
            source_type=source_type,
            **kwargs,
        )

    @abstractmethod
    def get_query_embedding(
        self, query: Any, source_type: Optional[EmbeddingSourceType] = "text", **kwargs
    ) -> list[float]:
        """
        Get embedding for a query (text or image).

        Args:
            query: Query text string or PIL Image object
            source_type: The type of source data ("text" or "image")

        Returns:
            List of float values representing the embedding
        """
        raise NotImplementedError()

    @abstractmethod
    def get_source_embedding(
        self, source: Any, source_type: Optional[EmbeddingSourceType] = "text", **kwargs
    ) -> list[float]:
        """
        Get embedding for a source field value (typically text).

        Args:
            source: Source field value (text)
            source_type: The type of source data ("text" or "image")

        Returns:
            List of float values representing the embedding
        """
        raise NotImplementedError()

    @abstractmethod
    def get_source_embeddings(
        self,
        sources: list[Any],
        source_type: Optional[EmbeddingSourceType] = "text",
        **kwargs,
    ) -> list[list[float]]:
        """
        Get embeddings for multiple source field values.

        Args:
            sources: List of source field values
            source_type: The type of source data ("text" or "image")

        Returns:
            List of embeddings, where each embedding is a list of float values
        """
        raise NotImplementedError()
