from __future__ import annotations

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

from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore

from langchain_community.vectorstores.utils import DistanceStrategy

logger = logging.getLogger(__name__)


class KDBAI(VectorStore):
    """`KDB.AI` vector store.

     See [https://kdb.ai](https://kdb.ai)

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

    Args:
        table: kdbai_client.Table object to use as storage,
        embedding: Any embedding function implementing
            `langchain.embeddings.base.Embeddings` interface,
        distance_strategy: One option from DistanceStrategy.EUCLIDEAN_DISTANCE,
            DistanceStrategy.DOT_PRODUCT or DistanceStrategy.COSINE.

    See the example [notebook](https://github.com/KxSystems/langchain/blob/KDB.AI/docs/docs/integrations/vectorstores/kdbai.ipynb).
    """

    def __init__(
        self,
        table: Any,
        embedding: Embeddings,
        distance_strategy: Optional[
            DistanceStrategy
        ] = DistanceStrategy.EUCLIDEAN_DISTANCE,
    ):
        try:
            import kdbai_client  # noqa
        except ImportError:
            raise ImportError(
                "Could not import kdbai_client python package. "
                "Please install it with `pip install kdbai_client`."
            )
        self._table = table
        self._embedding = embedding
        self.distance_strategy = distance_strategy

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

    def _embed_documents(self, texts: Iterable[str]) -> List[List[float]]:
        if isinstance(self._embedding, Embeddings):
            return self._embedding.embed_documents(list(texts))
        return [self._embedding(t) for t in texts]

    def _embed_query(self, text: str) -> List[float]:
        if isinstance(self._embedding, Embeddings):
            return self._embedding.embed_query(text)
        return self._embedding(text)

    def _insert(
        self,
        texts: List[str],
        ids: Optional[List[str]],
        metadata: Optional[Any] = None,
    ) -> None:
        try:
            import numpy as np
        except ImportError:
            raise ImportError(
                "Could not import numpy python package. "
                "Please install it with `pip install numpy`."
            )

        try:
            import pandas as pd
        except ImportError:
            raise ImportError(
                "Could not import pandas python package. "
                "Please install it with `pip install pandas`."
            )

        embeds = self._embedding.embed_documents(texts)
        df = pd.DataFrame()
        df["id"] = ids
        df["text"] = [t.encode("utf-8") for t in texts]
        df["embeddings"] = [np.array(e, dtype="float32") for e in embeds]
        if metadata is not None:
            df = pd.concat([df, metadata], axis=1)
        self._table.insert(df, warn=False)

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

        Args:
            texts (Iterable[str]): Texts to add to the vectorstore.
            metadatas (Optional[List[dict]]): List of metadata corresponding to each
                chunk of text.
            ids (Optional[List[str]]): List of IDs corresponding to each chunk of text.
            batch_size (Optional[int]): Size of batch of chunks of text to insert at
                once.

        Returns:
            List[str]: List of IDs of the added texts.
        """

        try:
            import pandas as pd
        except ImportError:
            raise ImportError(
                "Could not import pandas python package. "
                "Please install it with `pip install pandas`."
            )

        texts = list(texts)
        metadf: pd.DataFrame = None
        if metadatas is not None:
            if isinstance(metadatas, pd.DataFrame):
                metadf = metadatas
            else:
                metadf = pd.DataFrame(metadatas)
        out_ids: List[str] = []
        nbatches = (len(texts) - 1) // batch_size + 1
        for i in range(nbatches):
            istart = i * batch_size
            iend = (i + 1) * batch_size
            batch = texts[istart:iend]
            if ids:
                batch_ids = ids[istart:iend]
            else:
                batch_ids = [str(uuid.uuid4()) for _ in range(len(batch))]
            if metadf is not None:
                batch_meta = metadf.iloc[istart:iend].reset_index(drop=True)
            else:
                batch_meta = None
            self._insert(batch, batch_ids, batch_meta)
            out_ids = out_ids + batch_ids
        return out_ids

    def add_documents(
        self, documents: List[Document], batch_size: int = 32, **kwargs: Any
    ) -> List[str]:
        """Run more documents through the embeddings and add to the vectorstore.

        Args:
            documents (List[Document]: Documents to add to the vectorstore.
            batch_size (Optional[int]): Size of batch of documents to insert at once.

        Returns:
            List[str]: List of IDs of the added texts.
        """

        try:
            import pandas as pd
        except ImportError:
            raise ImportError(
                "Could not import pandas python package. "
                "Please install it with `pip install pandas`."
            )

        texts = [x.page_content for x in documents]
        metadata = pd.DataFrame([x.metadata for x in documents])
        return self.add_texts(texts, metadata=metadata, batch_size=batch_size)

    def similarity_search_with_score(
        self,
        query: str,
        k: int = 1,
        filter: Optional[List] = [],
        **kwargs: Any,
    ) -> List[Tuple[Document, float]]:
        """Run similarity search with distance from a query string.

        Args:
            query (str): Query string.
            k (Optional[int]): number of neighbors to retrieve.
            filter (Optional[List]): KDB.AI metadata filter clause: https://code.kx.com/kdbai/use/filter.html

        Returns:
            List[Document]: List of similar documents.
        """
        return self.similarity_search_by_vector_with_score(
            self._embed_query(query), k=k, filter=filter, **kwargs
        )

    def similarity_search_by_vector_with_score(
        self,
        embedding: List[float],
        *,
        k: int = 1,
        filter: Optional[List] = [],
        **kwargs: Any,
    ) -> List[Tuple[Document, float]]:
        """Return documents most similar to embedding, along with scores.

        Args:
            embedding (List[float]): query vector.
            k (Optional[int]): number of neighbors to retrieve.
            filter (Optional[List]): KDB.AI metadata filter clause: https://code.kx.com/kdbai/use/filter.html

        Returns:
            List[Document]: List of similar documents.
        """
        if "n" in kwargs:
            k = kwargs.pop("n")
        matches = self._table.search(vectors=[embedding], n=k, filter=filter, **kwargs)
        docs: list = []
        if isinstance(matches, list):
            matches = matches[0]
        else:
            return docs
        for row in matches.to_dict(orient="records"):
            text = row.pop("text")
            score = row.pop("__nn_distance")
            docs.append(
                (
                    Document(
                        page_content=text,
                        metadata={k: v for k, v in row.items() if k != "text"},
                    ),
                    score,
                )
            )
        return docs

    def similarity_search(
        self,
        query: str,
        k: int = 1,
        filter: Optional[List] = [],
        **kwargs: Any,
    ) -> List[Document]:
        """Run similarity search from a query string.

        Args:
            query (str): Query string.
            k (Optional[int]): number of neighbors to retrieve.
            filter (Optional[List]): KDB.AI metadata filter clause: https://code.kx.com/kdbai/use/filter.html

        Returns:
            List[Document]: List of similar documents.
        """
        docs_and_scores = self.similarity_search_with_score(
            query, k=k, filter=filter, **kwargs
        )
        return [doc for doc, _ in docs_and_scores]

    @classmethod
    def from_texts(
        cls: Any,
        texts: List[str],
        embedding: Embeddings,
        metadatas: Optional[List[dict]] = None,
        **kwargs: Any,
    ) -> Any:
        """Not implemented."""
        raise Exception("Not implemented.")