from __future__ import annotations

import logging
import uuid
from typing import Any, Iterable, List, Optional, Tuple

import numpy as np
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore

from langchain_community.vectorstores.utils import maximal_marginal_relevance

logger = logging.getLogger(__name__)


class Dingo(VectorStore):
    """`Dingo` vector store.

    To use, you should have the ``dingodb`` python package installed.

    Example:
        .. code-block:: python

            from langchain_community.vectorstores import Dingo
            from langchain_community.embeddings.openai import OpenAIEmbeddings

            embeddings = OpenAIEmbeddings()
            dingo = Dingo(embeddings, "text")
    """

    def __init__(
        self,
        embedding: Embeddings,
        text_key: str,
        *,
        client: Any = None,
        index_name: Optional[str] = None,
        dimension: int = 1024,
        host: Optional[List[str]] = None,
        user: str = "root",
        password: str = "123123",
        self_id: bool = False,
    ):
        """Initialize with Dingo client."""
        try:
            import dingodb
        except ImportError:
            raise ImportError(
                "Could not import dingo python package. "
                "Please install it with `pip install dingodb."
            )

        host = host if host is not None else ["172.20.31.10:13000"]

        # collection
        if client is not None:
            dingo_client = client
        else:
            try:
                # connect to dingo db
                dingo_client = dingodb.DingoDB(user, password, host)
            except ValueError as e:
                raise ValueError(f"Dingo failed to connect: {e}")

        self._text_key = text_key
        self._client = dingo_client

        if (
            index_name is not None
            and index_name not in dingo_client.get_index()
            and index_name.upper() not in dingo_client.get_index()
        ):
            if self_id is True:
                dingo_client.create_index(
                    index_name, dimension=dimension, auto_id=False
                )
            else:
                dingo_client.create_index(index_name, dimension=dimension)

        self._index_name = index_name
        self._embedding = embedding

    @property
    def embeddings(self) -> Optional[Embeddings]:
        return self._embedding

    def add_texts(
        self,
        texts: Iterable[str],
        metadatas: Optional[List[dict]] = None,
        ids: Optional[List[str]] = None,
        text_key: str = "text",
        batch_size: int = 500,
        **kwargs: Any,
    ) -> List[str]:
        """Run more texts through the embeddings and add to the vectorstore.

        Args:
            texts: Iterable of strings to add to the vectorstore.
            metadatas: Optional list of metadatas associated with the texts.
            ids: Optional list of ids to associate with the texts.

        Returns:
            List of ids from adding the texts into the vectorstore.

        """

        # Embed and create the documents
        ids = ids or [str(uuid.uuid4().int)[:13] for _ in texts]
        metadatas_list = []
        texts = list(texts)
        embeds = self._embedding.embed_documents(texts)
        for i, text in enumerate(texts):
            metadata = metadatas[i] if metadatas else {}
            metadata[self._text_key] = text
            metadatas_list.append(metadata)
        # upsert to Dingo
        for i in range(0, len(list(texts)), batch_size):
            j = i + batch_size
            add_res = self._client.vector_add(
                self._index_name, metadatas_list[i:j], embeds[i:j], ids[i:j]
            )
            if not add_res:
                raise Exception("vector add fail")

        return ids

    def similarity_search(
        self,
        query: str,
        k: int = 4,
        search_params: Optional[dict] = None,
        timeout: Optional[int] = None,
        **kwargs: Any,
    ) -> List[Document]:
        """Return Dingo documents most similar to query, along with scores.

        Args:
            query: Text to look up documents similar to.
            k: Number of Documents to return. Defaults to 4.
            search_params: Dictionary of argument(s) to filter on metadata

        Returns:
            List of Documents most similar to the query and score for each
        """
        docs_and_scores = self.similarity_search_with_score(
            query, k=k, search_params=search_params, **kwargs
        )
        return [doc for doc, _ in docs_and_scores]

    def similarity_search_with_score(
        self,
        query: str,
        k: int = 4,
        search_params: Optional[dict] = None,
        timeout: Optional[int] = None,
        **kwargs: Any,
    ) -> List[Tuple[Document, float]]:
        """Return Dingo documents most similar to query, along with scores.

        Args:
            query: Text to look up documents similar to.
            k: Number of Documents to return. Defaults to 4.
            search_params: Dictionary of argument(s) to filter on metadata

        Returns:
            List of Documents most similar to the query and score for each
        """
        docs = []
        query_obj = self._embedding.embed_query(query)
        results = self._client.vector_search(
            self._index_name, xq=query_obj, top_k=k, search_params=search_params
        )

        if not results:
            return []

        for res in results[0]["vectorWithDistances"]:
            score = res["distance"]
            if (
                "score_threshold" in kwargs
                and kwargs.get("score_threshold") is not None
            ):
                if score > kwargs.get("score_threshold"):
                    continue
            metadatas = res["scalarData"]
            id = res["id"]
            text = metadatas[self._text_key]["fields"][0]["data"]
            metadata = {"id": id, "text": text, "score": score}
            for meta_key in metadatas.keys():
                metadata[meta_key] = metadatas[meta_key]["fields"][0]["data"]
            docs.append((Document(page_content=text, metadata=metadata), score))

        return docs

    def max_marginal_relevance_search_by_vector(
        self,
        embedding: List[float],
        k: int = 4,
        fetch_k: int = 20,
        lambda_mult: float = 0.5,
        search_params: Optional[dict] = None,
        **kwargs: Any,
    ) -> List[Document]:
        """Return docs selected using the maximal marginal relevance.

        Maximal marginal relevance optimizes for similarity to query AND diversity
        among selected documents.

        Args:
            embedding: Embedding to look up documents similar to.
            k: Number of Documents to return. Defaults to 4.
            fetch_k: Number of Documents to fetch to pass to MMR algorithm.
            lambda_mult: Number between 0 and 1 that determines the degree
                        of diversity among the results with 0 corresponding
                        to maximum diversity and 1 to minimum diversity.
                        Defaults to 0.5.
        Returns:
            List of Documents selected by maximal marginal relevance.
        """
        results = self._client.vector_search(
            self._index_name, [embedding], search_params=search_params, top_k=k
        )

        mmr_selected = maximal_marginal_relevance(
            np.array([embedding], dtype=np.float32),
            [
                item["vector"]["floatValues"]
                for item in results[0]["vectorWithDistances"]
            ],
            k=k,
            lambda_mult=lambda_mult,
        )
        selected = []
        for i in mmr_selected:
            meta_data = {}
            for k, v in results[0]["vectorWithDistances"][i]["scalarData"].items():
                meta_data.update({str(k): v["fields"][0]["data"]})
            selected.append(meta_data)
        return [
            Document(page_content=metadata.pop(self._text_key), metadata=metadata)
            for metadata in selected
        ]

    def max_marginal_relevance_search(
        self,
        query: str,
        k: int = 4,
        fetch_k: int = 20,
        lambda_mult: float = 0.5,
        search_params: Optional[dict] = None,
        **kwargs: Any,
    ) -> List[Document]:
        """Return docs selected using the maximal marginal relevance.

        Maximal marginal relevance optimizes for similarity to query AND diversity
        among selected documents.

        Args:
            query: Text to look up documents similar to.
            k: Number of Documents to return. Defaults to 4.
            fetch_k: Number of Documents to fetch to pass to MMR algorithm.
            lambda_mult: Number between 0 and 1 that determines the degree
                        of diversity among the results with 0 corresponding
                        to maximum diversity and 1 to minimum diversity.
                        Defaults to 0.5.
        Returns:
            List of Documents selected by maximal marginal relevance.
        """
        embedding = self._embedding.embed_query(query)
        return self.max_marginal_relevance_search_by_vector(
            embedding, k, fetch_k, lambda_mult, search_params
        )

    @classmethod
    def from_texts(
        cls,
        texts: List[str],
        embedding: Embeddings,
        metadatas: Optional[List[dict]] = None,
        ids: Optional[List[str]] = None,
        text_key: str = "text",
        index_name: Optional[str] = None,
        dimension: int = 1024,
        client: Any = None,
        host: List[str] = ["172.20.31.10:13000"],
        user: str = "root",
        password: str = "123123",
        batch_size: int = 500,
        **kwargs: Any,
    ) -> Dingo:
        """Construct Dingo wrapper from raw documents.

                This is a user friendly interface that:
                    1. Embeds documents.
                    2. Adds the documents to a provided Dingo index

                This is intended to be a quick way to get started.

                Example:
                    .. code-block:: python

                        from langchain_community.vectorstores import Dingo
                        from langchain_community.embeddings import OpenAIEmbeddings
                        import dingodb
        sss
                        embeddings = OpenAIEmbeddings()
                        dingo = Dingo.from_texts(
                            texts,
                            embeddings,
                            index_name="langchain-demo"
                        )
        """
        try:
            import dingodb
        except ImportError:
            raise ImportError(
                "Could not import dingo python package. "
                "Please install it with `pip install dingodb`."
            )

        if client is not None:
            dingo_client = client
        else:
            try:
                # connect to dingo db
                dingo_client = dingodb.DingoDB(user, password, host)
            except ValueError as e:
                raise ValueError(f"Dingo failed to connect: {e}")
        if kwargs is not None and kwargs.get("self_id") is True:
            if (
                index_name is not None
                and index_name not in dingo_client.get_index()
                and index_name.upper() not in dingo_client.get_index()
            ):
                dingo_client.create_index(
                    index_name, dimension=dimension, auto_id=False
                )
        else:
            if (
                index_name is not None
                and index_name not in dingo_client.get_index()
                and index_name.upper() not in dingo_client.get_index()
            ):
                dingo_client.create_index(index_name, dimension=dimension)

        # Embed and create the documents

        ids = ids or [str(uuid.uuid4().int)[:13] for _ in texts]
        metadatas_list = []
        texts = list(texts)
        embeds = embedding.embed_documents(texts)
        for i, text in enumerate(texts):
            metadata = metadatas[i] if metadatas else {}
            metadata[text_key] = text
            metadatas_list.append(metadata)

        # upsert to Dingo
        for i in range(0, len(list(texts)), batch_size):
            j = i + batch_size
            add_res = dingo_client.vector_add(
                index_name, metadatas_list[i:j], embeds[i:j], ids[i:j]
            )
            if not add_res:
                raise Exception("vector add fail")
        return cls(embedding, text_key, client=dingo_client, index_name=index_name)

    def delete(
        self,
        ids: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> Any:
        """Delete by vector IDs or filter.
        Args:
            ids: List of ids to delete.
        """

        if ids is None:
            raise ValueError("No ids provided to delete.")

        return self._client.vector_delete(self._index_name, ids=ids)