# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
#
# Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
# SPDX-License-Identifier: MIT
import os
import re
import urllib.parse
from collections.abc import Callable
from typing import Any, Optional

from ....import_utils import optional_import_block, require_optional_import
from .base import Document, ItemID, QueryResults, VectorDB
from .utils import get_logger

with optional_import_block():
    import numpy as np
    import pgvector  # noqa: F401
    import psycopg
    from pgvector.psycopg import register_vector
    from sentence_transformers import SentenceTransformer

PGVECTOR_MAX_BATCH_SIZE = os.environ.get("PGVECTOR_MAX_BATCH_SIZE", 40000)
logger = get_logger(__name__)


@require_optional_import(["psycopg", "sentence_transformers", "numpy"], "retrievechat-pgvector")
class Collection:
    """A Collection object for PGVector.

    Attributes:
        client: The PGVector client.
        collection_name (str): The name of the collection. Default is "documents".
        embedding_function (Callable): The embedding function used to generate the vector representation.
            Default is None. SentenceTransformer("all-MiniLM-L6-v2").encode will be used when None.
            Models can be chosen from:
            https://huggingface.co/models?library=sentence-transformers
        metadata (Optional[dict[str, Any]]): The metadata of the collection.
        get_or_create (Optional): The flag indicating whether to get or create the collection.
    """

    def __init__(
        self,
        client: Any | None = None,
        collection_name: str = "ag2-docs",
        embedding_function: Callable[..., Any] | None = None,
        metadata: Any | None = None,
        get_or_create: Any | None = None,
    ):
        """Initialize the Collection object.

        Args:
            client: The PostgreSQL client.
            collection_name: The name of the collection. Default is "documents".
            embedding_function: The embedding function used to generate the vector representation.
            metadata: The metadata of the collection.
            get_or_create: The flag indicating whether to get or create the collection.

        Returns:
            None
        """
        self.client = client
        self.name = self.set_collection_name(collection_name)
        self.require_embeddings_or_documents = False
        self.ids = []
        if embedding_function:
            self.embedding_function = embedding_function
        else:
            self.embedding_function = SentenceTransformer("all-MiniLM-L6-v2").encode
        self.metadata = metadata if metadata else {"hnsw:space": "ip", "hnsw:construction_ef": 32, "hnsw:M": 16}
        self.documents = ""
        self.get_or_create = get_or_create
        # This will get the model dimension size by computing the embeddings dimensions
        sentences = [
            "The weather is lovely today in paradise.",
        ]
        embeddings = self.embedding_function(sentences)
        self.dimension = len(embeddings[0])

    def set_collection_name(self, collection_name) -> str:
        name = re.sub("-", "_", collection_name)
        self.name = name
        return self.name

    def add(
        self,
        ids: list[ItemID],
        documents: list[Document] | None,
        embeddings: list[Any] | None = None,
        metadatas: list[Any] | None = None,
    ) -> None:
        """Add documents to the collection.

        Args:
            ids (List[ItemID]): A list of document IDs.
            embeddings (List): A list of document embeddings. Optional
            metadatas (List): A list of document metadatas. Optional
            documents (List): A list of documents.

        Returns:
            None
        """
        cursor = self.client.cursor()
        sql_values = []
        if embeddings is not None and metadatas is not None:
            for doc_id, embedding, metadata, document in zip(ids, embeddings, metadatas, documents):
                metadata = re.sub("'", '"', str(metadata))
                sql_values.append((doc_id, embedding, metadata, document))
            sql_string = f"INSERT INTO {self.name} (id, embedding, metadatas, documents)\nVALUES (%s, %s, %s, %s);\n"
        elif embeddings is not None:
            for doc_id, embedding, document in zip(ids, embeddings, documents):
                sql_values.append((doc_id, embedding, document))
            sql_string = f"INSERT INTO {self.name} (id, embedding, documents) VALUES (%s, %s, %s);\n"
        elif metadatas is not None:
            for doc_id, metadata, document in zip(ids, metadatas, documents):
                metadata = re.sub("'", '"', str(metadata))
                embedding = self.embedding_function(document)
                sql_values.append((doc_id, metadata, embedding, document))
            sql_string = f"INSERT INTO {self.name} (id, metadatas, embedding, documents)\nVALUES (%s, %s, %s, %s);\n"
        else:
            for doc_id, document in zip(ids, documents):
                embedding = self.embedding_function(document)
                sql_values.append((doc_id, document, embedding))
            sql_string = f"INSERT INTO {self.name} (id, documents, embedding)\nVALUES (%s, %s, %s);\n"
        logger.debug(f"Add SQL String:\n{sql_string}\n{sql_values}")
        cursor.executemany(sql_string, sql_values)
        cursor.close()

    def upsert(
        self,
        ids: list[ItemID],
        documents: list[Document],
        embeddings: list[Any] | None = None,
        metadatas: list[Any] | None = None,
    ) -> None:
        """Upsert documents into the collection.

        Args:
            ids (List[ItemID]): A list of document IDs.
            documents (List): A list of documents.
            embeddings (List): A list of document embeddings.
            metadatas (List): A list of document metadatas.

        Returns:
            None
        """
        cursor = self.client.cursor()
        sql_values = []
        if embeddings is not None and metadatas is not None:
            for doc_id, embedding, metadata, document in zip(ids, embeddings, metadatas, documents):
                metadata = re.sub("'", '"', str(metadata))
                sql_values.append((doc_id, embedding, metadata, document, embedding, metadata, document))
            sql_string = (
                f"INSERT INTO {self.name} (id, embedding, metadatas, documents)\n"
                f"VALUES (%s, %s, %s, %s)\n"
                f"ON CONFLICT (id)\n"
                f"DO UPDATE SET embedding = %s,\n"
                f"metadatas = %s, documents = %s;\n"
            )
        elif embeddings is not None:
            for doc_id, embedding, document in zip(ids, embeddings, documents):
                sql_values.append((doc_id, embedding, document, embedding, document))
            sql_string = (
                f"INSERT INTO {self.name} (id, embedding, documents) "
                f"VALUES (%s, %s, %s) ON CONFLICT (id)\n"
                f"DO UPDATE SET embedding = %s, documents = %s;\n"
            )
        elif metadatas is not None:
            for doc_id, metadata, document in zip(ids, metadatas, documents):
                metadata = re.sub("'", '"', str(metadata))
                embedding = self.embedding_function(document)
                sql_values.append((doc_id, metadata, embedding, document, metadata, document, embedding))
            sql_string = (
                f"INSERT INTO {self.name} (id, metadatas, embedding, documents)\n"
                f"VALUES (%s, %s, %s, %s)\n"
                f"ON CONFLICT (id)\n"
                f"DO UPDATE SET metadatas = %s, documents = %s, embedding = %s;\n"
            )
        else:
            for doc_id, document in zip(ids, documents):
                embedding = self.embedding_function(document)
                sql_values.append((doc_id, document, embedding, document))
            sql_string = (
                f"INSERT INTO {self.name} (id, documents, embedding)\n"
                f"VALUES (%s, %s, %s)\n"
                f"ON CONFLICT (id)\n"
                f"DO UPDATE SET documents = %s;\n"
            )
        logger.debug(f"Upsert SQL String:\n{sql_string}\n{sql_values}")
        cursor.executemany(sql_string, sql_values)
        cursor.close()

    def count(self) -> int:
        """Get the total number of documents in the collection.

        Returns:
            int: The total number of documents.
        """
        cursor = self.client.cursor()
        query = f"SELECT COUNT(*) FROM {self.name}"
        cursor.execute(query)
        total = cursor.fetchone()[0]
        cursor.close()
        try:
            total = int(total)
        except (TypeError, ValueError):
            total = None
        return total

    def table_exists(self, table_name: str) -> bool:
        """Check if a table exists in the PostgreSQL database.

        Args:
            table_name (str): The name of the table to check.

        Returns:
            bool: True if the table exists, False otherwise.
        """
        cursor = self.client.cursor()
        cursor.execute(
            """
            SELECT EXISTS (
                SELECT 1
                FROM information_schema.tables
                WHERE table_name = %s
            )
            """,
            (table_name,),
        )
        exists = cursor.fetchone()[0]
        return exists

    def get(
        self,
        ids: str | None = None,
        include: str | None = None,
        where: str | None = None,
        limit: int | str | None = None,
        offset: int | str | None = None,
    ) -> list[Document]:
        """Retrieve documents from the collection.

        Args:
            ids (Optional[List]): A list of document IDs.
            include (Optional): The fields to include.
            where (Optional): Additional filtering criteria.
            limit (Optional): The maximum number of documents to retrieve.
            offset (Optional): The offset for pagination.

        Returns:
            List: The retrieved documents.
        """
        cursor = self.client.cursor()

        # Initialize variables for query components
        select_clause = "SELECT id, metadatas, documents, embedding"
        from_clause = f"FROM {self.name}"
        where_clause = ""
        limit_clause = ""
        offset_clause = ""

        # Handle include clause
        if include:
            select_clause = f"SELECT id, {', '.join(include)}, embedding"

        # Handle where clause
        if ids:
            where_clause = f"WHERE id IN ({', '.join(['%s' for _ in ids])})"
        elif where:
            where_clause = f"WHERE {where}"

        # Handle limit and offset clauses
        if limit:
            limit_clause = "LIMIT %s"
        if offset:
            offset_clause = "OFFSET %s"

        # Construct the full query
        query = f"{select_clause} {from_clause} {where_clause} {limit_clause} {offset_clause}"
        retrieved_documents = []
        try:
            # Execute the query with the appropriate values
            if ids is not None:
                cursor.execute(query, ids)
            else:
                query_params = []
                if limit:
                    query_params.append(limit)
                if offset:
                    query_params.append(offset)
                cursor.execute(query, query_params)

            retrieval = cursor.fetchall()
            for retrieved_document in retrieval:
                retrieved_documents.append(
                    Document(
                        id=retrieved_document[0].strip(),
                        metadata=retrieved_document[1],
                        content=retrieved_document[2],
                        embedding=retrieved_document[3],
                    )
                )
        except (psycopg.errors.UndefinedTable, psycopg.errors.UndefinedColumn) as e:
            logger.info(f"Error executing select on non-existent table: {self.name}. Creating it instead. Error: {e}")
            self.create_collection(collection_name=self.name, dimension=self.dimension)
            logger.info(f"Created table {self.name}")

        cursor.close()
        return retrieved_documents

    def update(self, ids: list[str], embeddings: list[Any], metadatas: list[Any], documents: list[Document]) -> None:
        """Update documents in the collection.

        Args:
            ids (List): A list of document IDs.
            embeddings (List): A list of document embeddings.
            metadatas (List): A list of document metadatas.
            documents (List): A list of documents.

        Returns:
            None
        """
        cursor = self.client.cursor()
        sql_values = []
        for doc_id, embedding, metadata, document in zip(ids, embeddings, metadatas, documents):
            sql_values.append((doc_id, embedding, metadata, document, doc_id, embedding, metadata, document))
        sql_string = (
            f"INSERT INTO {self.name} (id, embedding, metadata, document) "
            f"VALUES (%s, %s, %s, %s) "
            f"ON CONFLICT (id) "
            f"DO UPDATE SET id = %s, embedding = %s, "
            f"metadata = %s, document = %s;\n"
        )
        logger.debug(f"Upsert SQL String:\n{sql_string}\n")
        cursor.executemany(sql_string, sql_values)
        cursor.close()

    @staticmethod
    def euclidean_distance(arr1: list[float], arr2: list[float]) -> float:
        """Calculate the Euclidean distance between two vectors.

        Parameters:
        - arr1 (List[float]): The first vector.
        - arr2 (List[float]): The second vector.

        Returns:
        - float: The Euclidean distance between arr1 and arr2.
        """
        dist = np.linalg.norm(arr1 - arr2)
        return dist

    @staticmethod
    def cosine_distance(arr1: list[float], arr2: list[float]) -> float:
        """Calculate the cosine distance between two vectors.

        Parameters:
        - arr1 (List[float]): The first vector.
        - arr2 (List[float]): The second vector.

        Returns:
        - float: The cosine distance between arr1 and arr2.
        """
        dist = np.dot(arr1, arr2) / (np.linalg.norm(arr1) * np.linalg.norm(arr2))
        return dist

    @staticmethod
    def inner_product_distance(arr1: list[float], arr2: list[float]) -> float:
        """Calculate the Euclidean distance between two vectors.

        Parameters:
        - arr1 (List[float]): The first vector.
        - arr2 (List[float]): The second vector.

        Returns:
        - float: The Euclidean distance between arr1 and arr2.
        """
        dist = np.linalg.norm(arr1 - arr2)
        return dist

    def query(
        self,
        query_texts: list[str],
        collection_name: str | None = None,
        n_results: int | None = 10,
        distance_type: str | None = "euclidean",
        distance_threshold: float | None = -1,
        include_embedding: bool | None = False,
    ) -> QueryResults:
        """Query documents in the collection.

        Args:
            query_texts (List[str]): A list of query texts.
            collection_name (Optional[str]): The name of the collection.
            n_results (int): The maximum number of results to return.
            distance_type (Optional[str]): Distance search type - euclidean or cosine
            distance_threshold (Optional[float]): Distance threshold to limit searches
            include_embedding (Optional[bool]): Include embedding values in QueryResults
        Returns:
            QueryResults: The query results.
        """
        if collection_name:
            self.name = collection_name

        clause = "ORDER BY"
        if distance_threshold == -1:
            distance_threshold = ""
            clause = "ORDER BY"
        elif distance_threshold > 0:
            distance_threshold = f"< {distance_threshold}"
            clause = "WHERE"

        cursor = self.client.cursor()
        results = []
        for query_text in query_texts:
            vector = self.embedding_function(query_text, convert_to_tensor=False).tolist()
            if distance_type.lower() == "cosine":
                index_function = "<=>"
            elif distance_type.lower() == "euclidean":
                index_function = "<->"
            elif distance_type.lower() == "inner-product":
                index_function = "<#>"
            else:
                index_function = "<->"
            query = (
                f"SELECT id, documents, embedding, metadatas "
                f"FROM {self.name} "
                f"{clause} embedding {index_function} '{vector!s}' {distance_threshold} "
                f"LIMIT {n_results}"
            )
            cursor.execute(query)
            result = []
            for row in cursor.fetchall():
                fetched_document = Document(id=row[0].strip(), content=row[1], embedding=row[2], metadata=row[3])
                fetched_document_array = self.convert_string_to_array(array_string=fetched_document.get("embedding"))
                if distance_type.lower() == "cosine":
                    distance = self.cosine_distance(fetched_document_array, vector)
                elif distance_type.lower() == "euclidean":
                    distance = self.euclidean_distance(fetched_document_array, vector)
                elif distance_type.lower() == "inner-product":
                    distance = self.inner_product_distance(fetched_document_array, vector)
                else:
                    distance = self.euclidean_distance(fetched_document_array, vector)
                if not include_embedding:
                    fetched_document = Document(id=row[0].strip(), content=row[1], metadata=row[3])
                result.append((fetched_document, distance))
            results.append(result)
        cursor.close()
        logger.debug(f"Query Results: {results}")
        return results

    @staticmethod
    def convert_string_to_array(array_string: str) -> list[float]:
        """Convert a string representation of an array to a list of floats.

        Parameters:
        - array_string (str): The string representation of the array.

        Returns:
        - list: A list of floats parsed from the input string. If the input is
          not a string, it returns the input itself.
        """
        if not isinstance(array_string, str):
            return array_string
        array_string = array_string.strip("[]")
        array = [float(num) for num in array_string.split()]
        return array

    def modify(self, metadata, collection_name: str | None = None) -> None:
        """Modify metadata for the collection.

        Args:
            collection_name: The name of the collection.
            metadata: The new metadata.

        Returns:
            None
        """
        if collection_name:
            self.name = collection_name
        cursor = self.client.cursor()
        cursor.execute("UPDATE collectionsSET metadata = '%s'WHERE collection_name = '%s';", (metadata, self.name))
        cursor.close()

    def delete(self, ids: list[ItemID], collection_name: str | None = None) -> None:
        """Delete documents from the collection.

        Args:
            ids (List[ItemID]): A list of document IDs to delete.
            collection_name (str): The name of the collection to delete.

        Returns:
            None
        """
        if collection_name:
            self.name = collection_name
        cursor = self.client.cursor()
        id_placeholders = ", ".join(["%s" for _ in ids])
        cursor.execute(f"DELETE FROM {self.name} WHERE id IN ({id_placeholders});", ids)
        cursor.close()

    def delete_collection(self, collection_name: str | None = None) -> None:
        """Delete the entire collection.

        Args:
            collection_name (Optional[str]): The name of the collection to delete.

        Returns:
            None
        """
        if collection_name:
            self.name = collection_name
        cursor = self.client.cursor()
        cursor.execute(f"DROP TABLE IF EXISTS {self.name}")
        cursor.close()

    def create_collection(self, collection_name: str | None = None, dimension: str | int | None = None) -> None:
        """Create a new collection.

        Args:
            collection_name (Optional[str]): The name of the new collection.
            dimension (Optional[Union[str, int]]): The dimension size of the sentence embedding model

        Returns:
            None
        """
        if collection_name:
            self.name = collection_name

        if dimension:
            self.dimension = dimension
        elif self.dimension is None:
            self.dimension = 384

        cursor = self.client.cursor()
        cursor.execute(
            f"CREATE TABLE {self.name} ("
            f"documents text, id CHAR(8) PRIMARY KEY, metadatas JSONB, embedding vector({self.dimension}));"
            f"CREATE INDEX "
            f"ON {self.name} USING hnsw (embedding vector_l2_ops) WITH (m = {self.metadata['hnsw:M']}, "
            f"ef_construction = {self.metadata['hnsw:construction_ef']});"
            f"CREATE INDEX "
            f"ON {self.name} USING hnsw (embedding vector_cosine_ops) WITH (m = {self.metadata['hnsw:M']}, "
            f"ef_construction = {self.metadata['hnsw:construction_ef']});"
            f"CREATE INDEX "
            f"ON {self.name} USING hnsw (embedding vector_ip_ops) WITH (m = {self.metadata['hnsw:M']}, "
            f"ef_construction = {self.metadata['hnsw:construction_ef']});"
        )
        cursor.close()


@require_optional_import(["pgvector", "psycopg", "sentence_transformers"], "retrievechat-pgvector")
class PGVectorDB(VectorDB):
    """A vector database that uses PGVector as the backend."""

    def __init__(
        self,
        *,
        conn: Optional["psycopg.Connection"] = None,
        connection_string: str | None = None,
        host: str | None = None,
        port: int | str | None = None,
        dbname: str | None = None,
        username: str | None = None,
        password: str | None = None,
        connect_timeout: int | None = 10,
        embedding_function: Callable = None,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """Initialize the vector database.

        Note: connection_string or host + port + dbname must be specified

        Args:
            conn: psycopg.Connection | A customer connection object to connect to the database.
                A connection object may include additional key/values:
                https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING
            connection_string: "postgresql://username:password@hostname:port/database" | The PGVector connection string. Default is None.
            host: str | The host to connect to. Default is None.
            port: int | The port to connect to. Default is None.
            dbname: str | The database name to connect to. Default is None.
            username: str | The database username to use. Default is None.
            password: str | The database user password to use. Default is None.
            connect_timeout: int | The timeout to set for the connection. Default is 10.
            embedding_function: Callable | The embedding function used to generate the vector representation.
                Default is None. SentenceTransformer("all-MiniLM-L6-v2").encode will be used when None.
                Models can be chosen from:
                https://huggingface.co/models?library=sentence-transformers
            metadata: dict | The metadata of the vector database. Default is None. If None, it will use this
                setting: `{"hnsw:space": "ip", "hnsw:construction_ef": 30, "hnsw:M": 16}`. Creates Index on table
                using hnsw (embedding vector_l2_ops) WITH (m = hnsw:M) ef_construction = "hnsw:construction_ef".
                For more info: https://github.com/pgvector/pgvector?tab=readme-ov-file#hnsw
        Returns:
            None
        """
        self.client = self.establish_connection(
            conn=conn,
            connection_string=connection_string,
            host=host,
            port=port,
            dbname=dbname,
            username=username,
            password=password,
            connect_timeout=connect_timeout,
        )
        if embedding_function:
            self.embedding_function = embedding_function
        else:
            self.embedding_function = SentenceTransformer("all-MiniLM-L6-v2").encode
        self.metadata = metadata
        register_vector(self.client)
        self.active_collection = None

    def establish_connection(
        self,
        conn: Optional["psycopg.Connection"] = None,
        connection_string: str | None = None,
        host: str | None = None,
        port: int | str | None = None,
        dbname: str | None = None,
        username: str | None = None,
        password: str | None = None,
        connect_timeout: int | None = 10,
    ) -> "psycopg.Connection":
        """Establishes a connection to a PostgreSQL database using psycopg.

        Args:
            conn: An existing psycopg connection object. If provided, this connection will be used.
            connection_string: A string containing the connection information. If provided, a new connection will be established using this string.
            host: The hostname of the PostgreSQL server. Used if connection_string is not provided.
            port: The port number to connect to at the server host. Used if connection_string is not provided.
            dbname: The database name. Used if connection_string is not provided.
            username: The username to connect as. Used if connection_string is not provided.
            password: The user's password. Used if connection_string is not provided.
            connect_timeout: Maximum wait for connection, in seconds. The default is 10 seconds.

        Returns:
            A psycopg.Connection object representing the established connection.

        Raises:
            PermissionError if no credentials are supplied
            psycopg.Error: If an error occurs while trying to connect to the database.
        """
        try:
            if conn:
                self.client = conn
            elif connection_string:
                parsed_connection = urllib.parse.urlparse(connection_string)
                encoded_username = urllib.parse.quote(parsed_connection.username, safe="")
                encoded_password = urllib.parse.quote(parsed_connection.password, safe="")
                encoded_password = f":{encoded_password}@"
                encoded_host = urllib.parse.quote(parsed_connection.hostname, safe="")
                encoded_port = f":{parsed_connection.port}"
                encoded_database = urllib.parse.quote(parsed_connection.path[1:], safe="")
                connection_string_encoded = (
                    f"{parsed_connection.scheme}://{encoded_username}{encoded_password}"
                    f"{encoded_host}{encoded_port}/{encoded_database}"
                )
                self.client = psycopg.connect(conninfo=connection_string_encoded, autocommit=True)
            elif host:
                connection_string = ""
                if host:
                    encoded_host = urllib.parse.quote(host, safe="")
                    connection_string += f"host={encoded_host} "
                if port:
                    connection_string += f"port={port} "
                if dbname:
                    encoded_database = urllib.parse.quote(dbname, safe="")
                    connection_string += f"dbname={encoded_database} "
                if username:
                    encoded_username = urllib.parse.quote(username, safe="")
                    connection_string += f"user={encoded_username} "
                if password:
                    encoded_password = urllib.parse.quote(password, safe="")
                    connection_string += f"password={encoded_password} "

                self.client = psycopg.connect(
                    conninfo=connection_string,
                    connect_timeout=connect_timeout,
                    autocommit=True,
                )
            else:
                logger.error("Credentials were not supplied...")
                raise PermissionError
            self.client.execute("CREATE EXTENSION IF NOT EXISTS vector")
        except psycopg.Error as e:
            logger.error("Error connecting to the database: ", e)
            raise e
        return self.client

    def create_collection(
        self, collection_name: str, overwrite: bool = False, get_or_create: bool = True
    ) -> Collection:
        """Create a collection in the vector database.
        Case 1. if the collection does not exist, create the collection.
        Case 2. the collection exists, if overwrite is True, it will overwrite the collection.
        Case 3. the collection exists and overwrite is False, if get_or_create is True, it will get the collection,
            otherwise it raise a ValueError.

        Args:
            collection_name: str | The name of the collection.
            overwrite: bool | Whether to overwrite the collection if it exists. Default is False.
            get_or_create: bool | Whether to get the collection if it exists. Default is True.

        Returns:
            Collection | The collection object.
        """
        try:
            if self.active_collection and self.active_collection.name == collection_name:
                collection = self.active_collection
            else:
                collection = self.get_collection(collection_name)
        except ValueError:
            collection = None
        if collection is None:
            collection = Collection(
                client=self.client,
                collection_name=collection_name,
                embedding_function=self.embedding_function,
                get_or_create=get_or_create,
                metadata=self.metadata,
            )
            collection.set_collection_name(collection_name=collection_name)
            collection.create_collection(collection_name=collection_name)
            return collection
        elif overwrite:
            self.delete_collection(collection_name)
            collection = Collection(
                client=self.client,
                collection_name=collection_name,
                embedding_function=self.embedding_function,
                get_or_create=get_or_create,
                metadata=self.metadata,
            )
            collection.set_collection_name(collection_name=collection_name)
            collection.create_collection(collection_name=collection_name)
            return collection
        elif get_or_create:
            return collection
        elif not collection.table_exists(table_name=collection_name):
            collection = Collection(
                client=self.client,
                collection_name=collection_name,
                embedding_function=self.embedding_function,
                get_or_create=get_or_create,
                metadata=self.metadata,
            )
            collection.set_collection_name(collection_name=collection_name)
            collection.create_collection(collection_name=collection_name)
            return collection
        else:
            raise ValueError(f"Collection {collection_name} already exists.")

    def get_collection(self, collection_name: str = None) -> Collection:
        """Get the collection from the vector database.

        Args:
            collection_name: str | The name of the collection. Default is None. If None, return the
                current active collection.

        Returns:
            Collection | The collection object.
        """
        if collection_name is None:
            if self.active_collection is None:
                raise ValueError("No collection is specified.")
            else:
                logger.debug(
                    f"No collection is specified. Using current active collection {self.active_collection.name}."
                )
        else:
            if not (self.active_collection and self.active_collection.name == collection_name):
                self.active_collection = Collection(
                    client=self.client,
                    collection_name=collection_name,
                    embedding_function=self.embedding_function,
                )
        return self.active_collection

    def delete_collection(self, collection_name: str) -> None:
        """Delete the collection from the vector database.

        Args:
            collection_name: str | The name of the collection.

        Returns:
            None
        """
        if self.active_collection:
            self.active_collection.delete_collection(collection_name)
        else:
            collection = self.get_collection(collection_name)
            collection.delete_collection(collection_name)
        if self.active_collection and self.active_collection.name == collection_name:
            self.active_collection = None

    def _batch_insert(
        self, collection: Collection, embeddings=None, ids=None, metadatas=None, documents=None, upsert=False
    ) -> None:
        batch_size = int(PGVECTOR_MAX_BATCH_SIZE)
        default_metadata = {"hnsw:space": "ip", "hnsw:construction_ef": 32, "hnsw:M": 16}
        default_metadatas = [default_metadata] * min(batch_size, len(documents))
        for i in range(0, len(documents), min(batch_size, len(documents))):
            end_idx = i + min(batch_size, len(documents) - i)
            collection_kwargs = {
                "documents": documents[i:end_idx],
                "ids": ids[i:end_idx],
                "metadatas": metadatas[i:end_idx] if metadatas else default_metadatas,
                "embeddings": embeddings[i:end_idx] if embeddings else None,
            }
            if upsert:
                collection.upsert(**collection_kwargs)
            else:
                collection.add(**collection_kwargs)

    def insert_docs(self, docs: list[Document], collection_name: str = None, upsert: bool = False) -> None:
        """Insert documents into the collection of the vector database.

        Args:
            docs: List[Document] | A list of documents. Each document is a TypedDict `Document`.
            collection_name: str | The name of the collection. Default is None.
            upsert: bool | Whether to update the document if it exists. Default is False.
            kwargs: Dict | Additional keyword arguments.

        Returns:
            None
        """
        if not docs:
            return
        if docs[0].get("content") is None:
            raise ValueError("The document content is required.")
        if docs[0].get("id") is None:
            raise ValueError("The document id is required.")
        documents = [doc.get("content") for doc in docs]
        ids = [doc.get("id") for doc in docs]

        collection = self.get_collection(collection_name)
        if docs[0].get("embedding") is None:
            logger.debug(
                "No content embedding is provided. "
                "Will use the VectorDB's embedding function to generate the content embedding."
            )
            embeddings = None
        else:
            embeddings = [doc.get("embedding") for doc in docs]
        metadatas = None if docs[0].get("metadata") is None else [doc.get("metadata") for doc in docs]

        self._batch_insert(collection, embeddings, ids, metadatas, documents, upsert)

    def update_docs(self, docs: list[Document], collection_name: str = None) -> None:
        """Update documents in the collection of the vector database.

        Args:
            docs: List[Document] | A list of documents.
            collection_name: str | The name of the collection. Default is None.

        Returns:
            None
        """
        self.insert_docs(docs, collection_name, upsert=True)

    def delete_docs(self, ids: list[ItemID], collection_name: str = None) -> None:
        """Delete documents from the collection of the vector database.

        Args:
            ids: List[ItemID] | A list of document ids. Each id is a typed `ItemID`.
            collection_name: str | The name of the collection. Default is None.
            kwargs: Dict | Additional keyword arguments.

        Returns:
            None
        """
        collection = self.get_collection(collection_name)
        collection.delete(ids=ids, collection_name=collection_name)

    def retrieve_docs(
        self,
        queries: list[str],
        collection_name: str = None,
        n_results: int = 10,
        distance_threshold: float = -1,
    ) -> QueryResults:
        """Retrieve documents from the collection of the vector database based on the queries.

        Args:
            queries: List[str] | A list of queries. Each query is a string.
            collection_name: str | The name of the collection. Default is None.
            n_results: int | The number of relevant documents to return. Default is 10.
            distance_threshold: float | The threshold for the distance score, only distance smaller than it will be
                returned. Don't filter with it if `< 0`. Default is -1.
            kwargs: Dict | Additional keyword arguments.

        Returns:
            QueryResults | The query results. Each query result is a list of list of tuples containing the document and
                the distance.
        """
        collection = self.get_collection(collection_name)
        if isinstance(queries, str):
            queries = [queries]
        results = collection.query(
            query_texts=queries,
            n_results=n_results,
            distance_threshold=distance_threshold,
        )
        logger.debug(f"Retrieve Docs Results:\n{results}")
        return results

    def get_docs_by_ids(
        self, ids: list[ItemID] = None, collection_name: str = None, include=None, **kwargs
    ) -> list[Document]:
        """Retrieve documents from the collection of the vector database based on the ids.

        Args:
            ids: List[ItemID] | A list of document ids. If None, will return all the documents. Default is None.
            collection_name: str | The name of the collection. Default is None.
            include: List[str] | The fields to include. Default is None.
                If None, will include ["metadatas", "documents"], ids will always be included.
            kwargs: dict | Additional keyword arguments.

        Returns:
            List[Document] | The results.
        """
        collection = self.get_collection(collection_name)
        include = include if include else ["metadatas", "documents"]
        results = collection.get(ids, include=include, **kwargs)
        logger.debug(f"Retrieve Documents by ID Results:\n{results}")
        return results
