from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple

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

if TYPE_CHECKING:
    from transwarp_hippo_api.hippo_client import HippoClient

# Default connection
DEFAULT_HIPPO_CONNECTION = {
    "host": "localhost",
    "port": "7788",
    "username": "admin",
    "password": "admin",
}

logger = logging.getLogger(__name__)


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

    You need to install `hippo-api` and run Hippo.

    Please visit our official website for how to run a Hippo instance:
    https://www.transwarp.cn/starwarp

    Args:
        embedding_function (Embeddings): Function used to embed the text.
        table_name (str): Which Hippo table to use. Defaults to
            "test".
        database_name (str): Which Hippo database to use. Defaults to
            "default".
        number_of_shards (int): The number of shards for the Hippo table.Defaults to
            1.
        number_of_replicas (int): The number of replicas for the Hippo table.Defaults to
            1.
        connection_args (Optional[dict[str, any]]): The connection args used for
            this class comes in the form of a dict.
        index_params (Optional[dict]): Which index params to use. Defaults to
            IVF_FLAT.
        drop_old (Optional[bool]): Whether to drop the current collection. Defaults
            to False.
        primary_field (str): Name of the primary key field. Defaults to "pk".
        text_field (str): Name of the text field. Defaults to "text".
        vector_field (str): Name of the vector field. Defaults to "vector".

    The connection args used for this class comes in the form of a dict,
    here are a few of the options:
        host (str): The host of Hippo instance. Default at "localhost".
        port (str/int): The port of Hippo instance. Default at 7788.
        user (str): Use which user to connect to Hippo instance. If user and
            password are provided, we will add related header in every RPC call.
        password (str): Required when user is provided. The password
            corresponding to the user.

    Example:
        .. code-block:: python

        from langchain_community.vectorstores import Hippo
        from langchain_community.embeddings import OpenAIEmbeddings

        embedding = OpenAIEmbeddings()
        # Connect to a hippo instance on localhost
        vector_store = Hippo.from_documents(
            docs,
            embedding=embeddings,
            table_name="langchain_test",
            connection_args=HIPPO_CONNECTION
        )

    Raises:
        ValueError: If the hippo-api python package is not installed.
    """

    def __init__(
        self,
        embedding_function: Embeddings,
        table_name: str = "test",
        database_name: str = "default",
        number_of_shards: int = 1,
        number_of_replicas: int = 1,
        connection_args: Optional[Dict[str, Any]] = None,
        index_params: Optional[dict] = None,
        drop_old: Optional[bool] = False,
    ):
        self.number_of_shards = number_of_shards
        self.number_of_replicas = number_of_replicas
        self.embedding_func = embedding_function
        self.table_name = table_name
        self.database_name = database_name
        self.index_params = index_params

        # In order for a collection to be compatible,
        # 'pk' should be an auto-increment primary key and string
        self._primary_field = "pk"
        # In order for compatibility, the text field will need to be called "text"
        self._text_field = "text"
        # In order for compatibility, the vector field needs to be called "vector"
        self._vector_field = "vector"
        self.fields: List[str] = []
        # Create the connection to the server
        if connection_args is None:
            connection_args = DEFAULT_HIPPO_CONNECTION
        self.hc = self._create_connection_alias(connection_args)
        self.col: Any = None

        # If the collection exists, delete it
        try:
            if (
                self.hc.check_table_exists(self.table_name, self.database_name)
                and drop_old
            ):
                self.hc.delete_table(self.table_name, self.database_name)
        except Exception as e:
            logging.error(
                f"An error occurred while deleting the table " f"{self.table_name}: {e}"
            )
            raise

        try:
            if self.hc.check_table_exists(self.table_name, self.database_name):
                self.col = self.hc.get_table(self.table_name, self.database_name)
        except Exception as e:
            logging.error(
                f"An error occurred while getting the table " f"{self.table_name}: {e}"
            )
            raise

        # Initialize the vector database
        self._get_env()

    def _create_connection_alias(self, connection_args: dict) -> HippoClient:
        """Create the connection to the Hippo server."""
        # Grab the connection arguments that are used for checking existing connection
        try:
            from transwarp_hippo_api.hippo_client import HippoClient
        except ImportError as e:
            raise ImportError(
                "Unable to import transwarp_hipp_api, please install with "
                "`pip install hippo-api`."
            ) from e

        host: str = connection_args.get("host", None)
        port: int = connection_args.get("port", None)
        username: str = connection_args.get("username", "shiva")
        password: str = connection_args.get("password", "shiva")

        # Order of use is host/port, uri, address
        if host is not None and port is not None:
            if "," in host:
                hosts = host.split(",")
                given_address = ",".join([f"{h}:{port}" for h in hosts])
            else:
                given_address = str(host) + ":" + str(port)
        else:
            raise ValueError("Missing standard address type for reuse attempt")

        try:
            logger.info(f"create HippoClient[{given_address}]")
            return HippoClient([given_address], username=username, pwd=password)
        except Exception as e:
            logger.error("Failed to create new connection")
            raise e

    def _get_env(
        self, embeddings: Optional[list] = None, metadatas: Optional[List[dict]] = None
    ) -> None:
        logger.info("init ...")
        if embeddings is not None:
            logger.info("create collection")
            self._create_collection(embeddings, metadatas)
        self._extract_fields()
        self._create_index()

    def _create_collection(
        self, embeddings: list, metadatas: Optional[List[dict]] = None
    ) -> None:
        from transwarp_hippo_api.hippo_client import HippoField
        from transwarp_hippo_api.hippo_type import HippoType

        # Determine embedding dim
        dim = len(embeddings[0])
        logger.debug(f"[_create_collection] dim: {dim}")
        fields = []

        # Create the primary key field
        fields.append(HippoField(self._primary_field, True, HippoType.STRING))

        # Create the text field

        fields.append(HippoField(self._text_field, False, HippoType.STRING))

        # Create the vector field, supports binary or float vectors
        # to The binary vector type is to be developed.
        fields.append(
            HippoField(
                self._vector_field,
                False,
                HippoType.FLOAT_VECTOR,
                type_params={"dimension": dim},
            )
        )
        # to In Hippo,there is no method similar to the infer_type_data
        # types, so currently all non-vector data is converted to string type.

        if metadatas:
            #     # Create FieldSchema for each entry in metadata.
            for key, value in metadatas[0].items():
                #         # Infer the corresponding datatype of the metadata
                if isinstance(value, list):
                    value_dim = len(value)
                    fields.append(
                        HippoField(
                            key,
                            False,
                            HippoType.FLOAT_VECTOR,
                            type_params={"dimension": value_dim},
                        )
                    )
                else:
                    fields.append(HippoField(key, False, HippoType.STRING))

        logger.debug(f"[_create_collection] fields: {fields}")

        # Create the collection
        self.hc.create_table(
            name=self.table_name,
            auto_id=True,
            fields=fields,
            database_name=self.database_name,
            number_of_shards=self.number_of_shards,
            number_of_replicas=self.number_of_replicas,
        )
        self.col = self.hc.get_table(self.table_name, self.database_name)
        logger.info(
            f"[_create_collection] : "
            f"create table {self.table_name} in {self.database_name} successfully"
        )

    def _extract_fields(self) -> None:
        """Grab the existing fields from the Collection"""
        from transwarp_hippo_api.hippo_client import HippoTable

        if isinstance(self.col, HippoTable):
            schema = self.col.schema
            logger.debug(f"[_extract_fields] schema:{schema}")
            for x in schema:
                self.fields.append(x.name)
            logger.debug(f"04 [_extract_fields] fields:{self.fields}")

    # TO CAN: Translated into English, your statement would be: "Currently,
    # only the field named 'vector' (the automatically created vector field)
    # is checked for indexing. Indexes need to be created manually for other
    # vector type columns.
    def _get_index(self) -> Optional[Dict[str, Any]]:
        """Return the vector index information if it exists"""
        from transwarp_hippo_api.hippo_client import HippoTable

        if isinstance(self.col, HippoTable):
            table_info = self.hc.get_table_info(
                self.table_name, self.database_name
            ).get(self.table_name, {})
            embedding_indexes = table_info.get("embedding_indexes", None)
            if embedding_indexes is None:
                return None
            else:
                for x in self.hc.get_table_info(self.table_name, self.database_name)[
                    self.table_name
                ]["embedding_indexes"]:
                    logger.debug(f"[_get_index] embedding_indexes {embedding_indexes}")
                    if x["column"] == self._vector_field:
                        return x
        return None

    # TO Indexes can only be created for the self._vector_field field.
    def _create_index(self) -> None:
        """Create a index on the collection"""
        from transwarp_hippo_api.hippo_client import HippoTable
        from transwarp_hippo_api.hippo_type import IndexType, MetricType

        if isinstance(self.col, HippoTable) and self._get_index() is None:
            if self._get_index() is None:
                if self.index_params is None:
                    self.index_params = {
                        "index_name": "langchain_auto_create",
                        "metric_type": MetricType.L2,
                        "index_type": IndexType.IVF_FLAT,
                        "nlist": 10,
                    }

                    self.col.create_index(
                        self._vector_field,
                        self.index_params["index_name"],
                        self.index_params["index_type"],
                        self.index_params["metric_type"],
                        nlist=self.index_params["nlist"],
                    )
                    logger.debug(
                        self.col.activate_index(self.index_params["index_name"])
                    )
                    logger.info("create index successfully")
                else:
                    index_dict = {
                        "IVF_FLAT": IndexType.IVF_FLAT,
                        "FLAT": IndexType.FLAT,
                        "IVF_SQ": IndexType.IVF_SQ,
                        "IVF_PQ": IndexType.IVF_PQ,
                        "HNSW": IndexType.HNSW,
                    }

                    metric_dict = {
                        "ip": MetricType.IP,
                        "IP": MetricType.IP,
                        "l2": MetricType.L2,
                        "L2": MetricType.L2,
                    }
                    self.index_params["metric_type"] = metric_dict[
                        self.index_params["metric_type"]
                    ]

                    if self.index_params["index_type"] == "FLAT":
                        self.index_params["index_type"] = index_dict[
                            self.index_params["index_type"]
                        ]
                        self.col.create_index(
                            self._vector_field,
                            self.index_params["index_name"],
                            self.index_params["index_type"],
                            self.index_params["metric_type"],
                        )
                        logger.debug(
                            self.col.activate_index(self.index_params["index_name"])
                        )
                    elif (
                        self.index_params["index_type"] == "IVF_FLAT"
                        or self.index_params["index_type"] == "IVF_SQ"
                    ):
                        self.index_params["index_type"] = index_dict[
                            self.index_params["index_type"]
                        ]
                        self.col.create_index(
                            self._vector_field,
                            self.index_params["index_name"],
                            self.index_params["index_type"],
                            self.index_params["metric_type"],
                            nlist=self.index_params.get("nlist", 10),
                            nprobe=self.index_params.get("nprobe", 10),
                        )
                        logger.debug(
                            self.col.activate_index(self.index_params["index_name"])
                        )
                    elif self.index_params["index_type"] == "IVF_PQ":
                        self.index_params["index_type"] = index_dict[
                            self.index_params["index_type"]
                        ]
                        self.col.create_index(
                            self._vector_field,
                            self.index_params["index_name"],
                            self.index_params["index_type"],
                            self.index_params["metric_type"],
                            nlist=self.index_params.get("nlist", 10),
                            nprobe=self.index_params.get("nprobe", 10),
                            nbits=self.index_params.get("nbits", 8),
                            m=self.index_params.get("m"),
                        )
                        logger.debug(
                            self.col.activate_index(self.index_params["index_name"])
                        )
                    elif self.index_params["index_type"] == "HNSW":
                        self.index_params["index_type"] = index_dict[
                            self.index_params["index_type"]
                        ]
                        self.col.create_index(
                            self._vector_field,
                            self.index_params["index_name"],
                            self.index_params["index_type"],
                            self.index_params["metric_type"],
                            M=self.index_params.get("M"),
                            ef_construction=self.index_params.get("ef_construction"),
                            ef_search=self.index_params.get("ef_search"),
                        )
                        logger.debug(
                            self.col.activate_index(self.index_params["index_name"])
                        )
                    else:
                        raise ValueError(
                            "Index name does not match, "
                            "please enter the correct index name. "
                            "(FLAT, IVF_FLAT, IVF_PQ,IVF_SQ, HNSW)"
                        )

    def add_texts(
        self,
        texts: Iterable[str],
        metadatas: Optional[List[dict]] = None,
        timeout: Optional[int] = None,
        batch_size: int = 1000,
        **kwargs: Any,
    ) -> List[str]:
        """
        Add text to the collection.

        Args:
            texts: An iterable that contains the text to be added.
            metadatas: An optional list of dictionaries,
            each dictionary contains the metadata associated with a text.
            timeout: Optional timeout, in seconds.
            batch_size: The number of texts inserted in each batch, defaults to 1000.
            **kwargs: Other optional parameters.

        Returns:
            A list of strings, containing the unique identifiers of the inserted texts.

        Note:
            If the collection has not yet been created,
            this method will create a new collection.
        """
        from transwarp_hippo_api.hippo_client import HippoTable

        if not texts or all(t == "" for t in texts):
            logger.debug("Nothing to insert, skipping.")
            return []
        texts = list(texts)

        logger.debug(f"[add_texts] texts: {texts}")

        try:
            embeddings = self.embedding_func.embed_documents(texts)
        except NotImplementedError:
            embeddings = [self.embedding_func.embed_query(x) for x in texts]

        if len(embeddings) == 0:
            logger.debug("Nothing to insert, skipping.")
            return []

        logger.debug(f"[add_texts] len_embeddings:{len(embeddings)}")

        # 如果还没有创建collection则创建collection
        if not isinstance(self.col, HippoTable):
            self._get_env(embeddings, metadatas)

        # Dict to hold all insert columns
        insert_dict: Dict[str, list] = {
            self._text_field: texts,
            self._vector_field: embeddings,
        }
        logger.debug(f"[add_texts] metadatas:{metadatas}")
        logger.debug(f"[add_texts] fields:{self.fields}")
        if metadatas is not None:
            for d in metadatas:
                for key, value in d.items():
                    if key in self.fields:
                        insert_dict.setdefault(key, []).append(value)

        logger.debug(insert_dict[self._text_field])

        # Total insert count
        vectors: list = insert_dict[self._vector_field]
        total_count = len(vectors)

        if "pk" in self.fields:
            self.fields.remove("pk")

        logger.debug(f"[add_texts] total_count:{total_count}")
        for i in range(0, total_count, batch_size):
            # Grab end index
            end = min(i + batch_size, total_count)
            # Convert dict to list of lists batch for insertion
            insert_list = [insert_dict[x][i:end] for x in self.fields]
            try:
                res = self.col.insert_rows(insert_list)
                logger.info(f"05 [add_texts] insert {res}")
            except Exception as e:
                logger.error(
                    "Failed to insert batch starting at entity: %s/%s", i, total_count
                )
                raise e
        return [""]

    def similarity_search(
        self,
        query: str,
        k: int = 4,
        param: Optional[dict] = None,
        expr: Optional[str] = None,
        timeout: Optional[int] = None,
        **kwargs: Any,
    ) -> List[Document]:
        """
        Perform a similarity search on the query string.

        Args:
            query (str): The text to search for.
            k (int, optional): The number of results to return. Default is 4.
            param (dict, optional): Specifies the search parameters for the index.
            Defaults to None.
            expr (str, optional): Filtering expression. Defaults to None.
            timeout (int, optional): Time to wait before a timeout error.
            Defaults to None.
            kwargs: Keyword arguments for Collection.search().

        Returns:
            List[Document]: The document results of the search.
        """

        if self.col is None:
            logger.debug("No existing collection to search.")
            return []
        res = self.similarity_search_with_score(
            query=query, k=k, param=param, expr=expr, timeout=timeout, **kwargs
        )
        return [doc for doc, _ in res]

    def similarity_search_with_score(
        self,
        query: str,
        k: int = 4,
        param: Optional[dict] = None,
        expr: Optional[str] = None,
        timeout: Optional[int] = None,
        **kwargs: Any,
    ) -> List[Tuple[Document, float]]:
        """
        Performs a search on the query string and returns results with scores.

        Args:
            query (str): The text being searched.
            k (int, optional): The number of results to return.
            Default is 4.
            param (dict): Specifies the search parameters for the index.
            Default is None.
            expr (str, optional): Filtering expression. Default is None.
            timeout (int, optional): The waiting time before a timeout error.
            Default is None.
            kwargs: Keyword arguments for Collection.search().

        Returns:
            List[float], List[Tuple[Document, any, any]]:
        """

        if self.col is None:
            logger.debug("No existing collection to search.")
            return []

        # Embed the query text.
        embedding = self.embedding_func.embed_query(query)

        ret = self.similarity_search_with_score_by_vector(
            embedding=embedding, k=k, param=param, expr=expr, timeout=timeout, **kwargs
        )
        return ret

    def similarity_search_with_score_by_vector(
        self,
        embedding: List[float],
        k: int = 4,
        param: Optional[dict] = None,
        expr: Optional[str] = None,
        timeout: Optional[int] = None,
        **kwargs: Any,
    ) -> List[Tuple[Document, float]]:
        """
        Performs a search on the query string and returns results with scores.

        Args:
            embedding (List[float]): The embedding vector being searched.
            k (int, optional): The number of results to return.
            Default is 4.
            param (dict): Specifies the search parameters for the index.
            Default is None.
            expr (str, optional): Filtering expression. Default is None.
            timeout (int, optional): The waiting time before a timeout error.
            Default is None.
            kwargs: Keyword arguments for Collection.search().

        Returns:
            List[Tuple[Document, float]]: Resulting documents and scores.
        """
        if self.col is None:
            logger.debug("No existing collection to search.")
            return []

        # if param is None:
        #     param = self.search_params

        # Determine result metadata fields.
        output_fields = self.fields[:]
        output_fields.remove(self._vector_field)

        # Perform the search.
        logger.debug(f"search_field:{self._vector_field}")
        logger.debug(f"vectors:{[embedding]}")
        logger.debug(f"output_fields:{output_fields}")
        logger.debug(f"topk:{k}")
        logger.debug(f"dsl:{expr}")

        res = self.col.query(
            search_field=self._vector_field,
            vectors=[embedding],
            output_fields=output_fields,
            topk=k,
            dsl=expr,
        )
        # Organize results.
        logger.debug(f"[similarity_search_with_score_by_vector] res:{res}")
        score_col = self._text_field + "%scores"
        ret = []
        count = 0
        for items in zip(*[res[0][field] for field in output_fields]):
            meta = {field: value for field, value in zip(output_fields, items)}
            doc = Document(page_content=meta.pop(self._text_field), metadata=meta)
            logger.debug(
                f"[similarity_search_with_score_by_vector] "
                f"res[0][score_col]:{res[0][score_col]}"
            )
            score = res[0][score_col][count]
            count += 1
            ret.append((doc, score))

        return ret

    @classmethod
    def from_texts(
        cls,
        texts: List[str],
        embedding: Embeddings,
        metadatas: Optional[List[dict]] = None,
        table_name: str = "test",
        database_name: str = "default",
        connection_args: Dict[str, Any] = DEFAULT_HIPPO_CONNECTION,
        index_params: Optional[Dict[Any, Any]] = None,
        search_params: Optional[Dict[str, Any]] = None,
        drop_old: bool = False,
        **kwargs: Any,
    ) -> "Hippo":
        """
        Creates an instance of the VST class from the given texts.

        Args:
            texts (List[str]): List of texts to be added.
            embedding (Embeddings): Embedding model for the texts.
            metadatas (List[dict], optional):
            List of metadata dictionaries for each text.Defaults to None.
            table_name (str): Name of the table. Defaults to "test".
            database_name (str): Name of the database. Defaults to "default".
            connection_args (dict[str, Any]): Connection parameters.
            Defaults to DEFAULT_HIPPO_CONNECTION.
            index_params (dict): Indexing parameters. Defaults to None.
            search_params (dict): Search parameters. Defaults to an empty dictionary.
            drop_old (bool): Whether to drop the old collection. Defaults to False.
            kwargs: Other arguments.

        Returns:
            Hippo: An instance of the VST class.
        """

        if search_params is None:
            search_params = {}
        logger.info("00 [from_texts] init the class of Hippo")
        vector_db = cls(
            embedding_function=embedding,
            table_name=table_name,
            database_name=database_name,
            connection_args=connection_args,
            index_params=index_params,
            drop_old=drop_old,
            **kwargs,
        )
        logger.debug(f"[from_texts] texts:{texts}")
        logger.debug(f"[from_texts] metadatas:{metadatas}")
        vector_db.add_texts(texts=texts, metadatas=metadatas)
        return vector_db