"""Default vector database base with multi-modal embedding support.

This module provides ``DefaultVectorDatabase``, an abstract class that
combines ``AugmentationDefault`` (query normalization) with
``BaseVectorDatabase`` (storage interface) and adds a default multi-modal
embedding function. The embedding function dispatches to the appropriate
API method (text, image, audio, video, document, table) based on the
media type of the informations being embedded, so that concrete
implementations only need to handle storage and retrieval.
"""

import json
from typing import Callable, List, Optional, Sequence

from pydantic import BaseModel

from litemind.agent.augmentations.augmentation_default import AugmentationDefault
from litemind.agent.augmentations.information.information import Information
from litemind.agent.augmentations.vector_db.base_vector_db import BaseVectorDatabase
from litemind.apis.base_api import BaseApi
from litemind.media.types.media_audio import Audio
from litemind.media.types.media_code import Code
from litemind.media.types.media_document import Document
from litemind.media.types.media_image import Image
from litemind.media.types.media_json import Json
from litemind.media.types.media_object import Object
from litemind.media.types.media_table import Table
from litemind.media.types.media_text import Text
from litemind.media.types.media_video import Video
from litemind.utils.uri_utils import is_uri


class DefaultVectorDatabase(AugmentationDefault, BaseVectorDatabase):
    """Abstract base for vector databases with a default embedding function.

    Provides a multi-modal embedding function that dispatches to the
    appropriate API method (text, image, audio, etc.) based on the media
    type of the informations being embedded.
    """

    def __init__(
        self,
        name: Optional[str] = None,
        description: Optional[str] = None,
        embedding_function: Optional[
            Callable[[List[Information]], List[List[float]]]
        ] = None,
        embedding_length: Optional[int] = None,
        api: BaseApi = None,
    ):
        """
        Initialize the vector database with an embedding function.

        Parameters
        ----------
        name : str, optional
            The name of the augmentation. Defaults to the class name.
        description : Optional[str]
            A description of the augmentation.
        embedding_function : Optional[Callable[[List[Information]], List[List[float]]]]
            A function that takes a list of informations and returns
            embedding vectors. If not provided, uses the built-in
            multi-modal embedding function.
        embedding_length : Optional[int]
            The dimensionality of the embedding vectors. If not provided,
            it is auto-detected from a sample embedding.
        api : BaseApi, optional
            The API to use for embedding. If not provided, uses
            ``CombinedApi()`` as the default.
        """
        super().__init__(name=name, description=description)

        # Set the name of the augmentation to the extended class name:
        if name is None:
            name = self.__class__.__name__
        self.name = name

        # Set the embedding function:
        self.embedding_function = embedding_function or self._default_embedding_function

        # If no api is provided, use the default API (only needed for default embedding):
        if api is None and embedding_function is None:
            try:
                from litemind import CombinedApi

                api = CombinedApi()
            except ValueError:
                raise ValueError(
                    "No API provided and litemind's default API is not available. "
                    "Please provide an API."
                )
        self.api = api

        # Call the embedding function to get the embedding length if not loaded from disk
        if embedding_length is None:
            # Make a simple information that contains some text:
            query_info = Information(Text("sample text"))

            sample_embedding = self.embedding_function([query_info])[0]
            embedding_length = len(sample_embedding)

        # Set the embedding length:
        self.embedding_length = embedding_length

    def _default_embedding_function(
        self, informations: Sequence[Information]
    ) -> Sequence[Sequence[float]]:
        """
        Default embedding function using litemind's API.

        Parameters
        ----------
        informations: Sequence[Information]
            List of informations to embed.

        Returns
        -------
        List[List[float]]
            The embedding vectors.
        """

        # If embedding_length does not exist in object, then we use a default value for the embedding length:
        if not hasattr(self, "embedding_length"):
            self.embedding_length = 768

        # First we check if all the informations are of the same type:
        info_types = {info.type for info in informations}
        if len(info_types) > 1:

            # Initialise the list of embeddings:
            embeddings = []

            # If the informations are of different types we must handle them separately:
            for info in informations:
                # Compute the embedding for the information:
                embeddings_for_info = self._raw_embedding_function([info])

                # Add embedding to the list of embeddings:
                embeddings.append(embeddings_for_info[0])

        else:
            # If all the informations are of the same type, we can proceed normally:
            embeddings = self._raw_embedding_function(informations)

        # Return the embeddings:
        return embeddings

    def _raw_embedding_function(
        self, informations: Sequence[Information]
    ) -> List[List[float]]:
        """
        Compute embeddings for a batch of same-type informations.

        Dispatches to the appropriate API embedding method based on
        media type (text, image, audio, video, document, etc.).

        Parameters
        ----------
        informations : Sequence[Information]
            Informations to embed. All must have the same media type.

        Returns
        -------
        List[List[float]]
            The embedding vectors.

        Raises
        ------
        ValueError
            If the list is empty or the media type is unsupported.
        """

        # Check if list is empty:
        if not informations:
            raise ValueError("No informations provided for embedding.")

        # If the information already has an embedding, we use it:
        if all(info.embedding is not None for info in informations):
            embeddings = [info.embedding for info in informations]
        else:

            if all(info.has_type(Text) for info in informations):

                # If the information is of type text, we use the text embedding function:
                texts = [info.content for info in informations]

                # Embed the texts:
                embeddings = self.api.embed_texts(
                    texts=texts, dimensions=self.embedding_length
                )
            elif all(info.has_type(Code) for info in informations):
                # If the information is of type code, we use the code embedding function:
                codes = [info.content for info in informations]

                # Embed the codes:
                embeddings = self.api.embed_texts(
                    texts=codes, dimensions=self.embedding_length
                )

            elif all(info.has_type(Json) for info in informations):
                # If the information is of type json, we use the json embedding function:
                json_strs = [
                    (
                        json.dumps(info.content)
                        if isinstance(info.content, dict)
                        else info.content
                    )
                    for info in informations
                ]

                # Embed the json:
                embeddings = self.api.embed_texts(
                    texts=json_strs, dimensions=self.embedding_length
                )

            elif all(info.has_type(Object) for info in informations):
                # If the information is of type object, we use the object embedding function:

                # First let's check that the object is a Pydantic BaseModel:
                if any(
                    not isinstance(info.content, BaseModel) for info in informations
                ):
                    raise ValueError(
                        "Information content must be a Pydantic BaseModel for object type Object."
                    )

                # Convert to JSON strings:
                json_strs = [
                    json.dumps(info.content.model_dump(mode="json"))
                    for info in informations
                ]

                # Embed the objects:
                embeddings = self.api.embed_texts(
                    texts=json_strs, dimensions=self.embedding_length
                )

            elif all(info.has_type(Image) for info in informations):
                # If the information is of type image, we use the image embedding function:
                image_uris = [info.content for info in informations]

                # Embed the images:
                embeddings = self.api.embed_images(
                    image_uris=image_uris, dimensions=self.embedding_length
                )
            elif all(info.has_type(Audio) for info in informations):
                # If the information is of type audio, we use the audio embedding function:
                audio_uris = [info.content for info in informations]

                # Embed the audio files:
                embeddings = self.api.embed_audios(
                    audio_uris=audio_uris, dimensions=self.embedding_length
                )
            elif all(info.has_type(Video) for info in informations):
                # If the information is of type video, we use the video embedding function:
                video_uris = [info.content for info in informations]

                # Embed the videos:
                embeddings = self.api.embed_videos(
                    video_uris=video_uris, dimensions=self.embedding_length
                )

            elif all(info.has_type(Document) for info in informations):
                # If the information is of type document, we use the document embedding function:
                document_uris = [info.content for info in informations]

                # Embed the videos:
                embeddings = self.api.embed_documents(
                    document_uris=document_uris, dimensions=self.embedding_length
                )

            elif all(info.has_type(Table) for info in informations):

                # First let's check that the all informations have a URIs:
                if any(not is_uri(info.content) for info in informations):
                    raise ValueError(
                        "All informations content must be a URI for table type."
                    )

                # Convert to Markdown:
                table_markdowns = [info.media.to_markdown() for info in informations]

                # Embed the table:
                embeddings = self.api.embed_texts(
                    texts=table_markdowns, dimensions=self.embedding_length
                )

            else:
                # If the information is of an unknown type, we raise an error:
                raise ValueError(f"Unknown information type: {informations[0].type}.")

        return embeddings

    def __del__(self):
        """Ensure the connection is closed when the object is garbage collected."""
        try:
            self.close()
        except:
            pass
