# 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 json
import time
from datetime import timedelta
from typing import Any, Callable, Literal, 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
    from couchbase import search
    from couchbase.auth import PasswordAuthenticator
    from couchbase.cluster import Cluster, ClusterOptions
    from couchbase.collection import Collection
    from couchbase.management.search import SearchIndex
    from couchbase.options import SearchOptions
    from couchbase.vector_search import VectorQuery, VectorSearch
    from sentence_transformers import SentenceTransformer


logger = get_logger(__name__)

DEFAULT_BATCH_SIZE = 1000
_SAMPLE_SENTENCE = ["The weather is lovely today in paradise."]
TEXT_KEY = "content"
EMBEDDING_KEY = "embedding"


@require_optional_import(["couchbase", "sentence_transformers"], "retrievechat-couchbase")
class CouchbaseVectorDB(VectorDB):
    """
    A vector database implementation that uses Couchbase as the backend.
    """

    def __init__(
        self,
        connection_string: str = "couchbase://localhost",
        username: str = "Administrator",
        password: str = "password",
        bucket_name: str = "vector_db",
        embedding_function: Callable = None,
        scope_name: str = "_default",
        collection_name: str = "_default",
        index_name: str = None,
    ):
        """
        Initialize the vector database.
        Args:
            connection_string (str): The Couchbase connection string to connect to. Default is 'couchbase://localhost'.
            username (str): The username for Couchbase authentication. Default is 'Administrator'.
            password (str): The password for Couchbase authentication. Default is 'password'.
            bucket_name (str): The name of the bucket. Default is 'vector_db'.
            embedding_function (Callable): The embedding function used to generate the vector representation. Default is SentenceTransformer("all-MiniLM-L6-v2").encode.
            scope_name (str): The name of the scope. Default is '_default'.
            collection_name (str): The name of the collection to create for this vector database. Default is '_default'.
            index_name (str): Index name for the vector database. Default is None.
            overwrite (bool): Whether to overwrite existing data. Default is False.
            wait_until_index_ready (float or None): Blocking call to wait until the database indexes are ready. None means no wait. Default is None.
            wait_until_document_ready (float or None): Blocking call to wait until the database documents are ready. None means no wait. Default is None.
        """
        if embedding_function is None:
            embedding_function = SentenceTransformer("all-MiniLM-L6-v2").encode
        self.embedding_function = embedding_function
        self.index_name = index_name

        # This will get the model dimension size by computing the embeddings dimensions
        self.dimensions = self._get_embedding_size()

        try:
            auth = PasswordAuthenticator(username, password)
            cluster = Cluster(connection_string, ClusterOptions(auth))
            cluster.wait_until_ready(timedelta(seconds=5))
            self.cluster = cluster

            self.bucket = cluster.bucket(bucket_name)
            self.scope = self.bucket.scope(scope_name)
            self.collection = self.scope.collection(collection_name)
            self.active_collection = self.collection

            logger.debug("Successfully connected to Couchbase")
        except Exception as err:
            raise ConnectionError("Could not connect to Couchbase server") from err

    def search_index_exists(self, index_name: str):
        """Check if the specified index is ready"""
        try:
            search_index_mgr = self.scope.search_indexes()
            index = search_index_mgr.get_index(index_name)
            return index.is_valid()
        except Exception:
            return False

    def _get_embedding_size(self):
        return len(self.embedding_function(_SAMPLE_SENTENCE)[0])

    def create_collection(
        self,
        collection_name: str,
        overwrite: bool = False,
        get_or_create: bool = True,
    ) -> "Collection":
        """
        Create a collection in the vector database and create a vector search index in the collection.
        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 or create the collection. Default is True
        """
        if overwrite:
            self.delete_collection(collection_name)

        try:
            collection_mgr = self.bucket.collections()
            collection_mgr.create_collection(self.scope.name, collection_name)
            self.cluster.query(f"CREATE PRIMARY INDEX ON {self.bucket.name}.{self.scope.name}.{collection_name}")

        except Exception:
            if not get_or_create:
                raise ValueError(f"Collection {collection_name} already exists.")
            else:
                logger.debug(f"Collection {collection_name} already exists. Getting the collection.")

        collection = self.scope.collection(collection_name)
        self.create_index_if_not_exists(index_name=self.index_name, collection=collection)
        return collection

    def create_index_if_not_exists(
        self, index_name: str = "vector_index", collection: Optional["Collection"] = None
    ) -> None:
        """
        Creates a vector search index on the specified collection in Couchbase.
        Args:
            index_name (str, optional): The name of the vector search index to create. Defaults to "vector_search_index".
            collection (Collection, optional): The Couchbase collection to create the index on. Defaults to None.
        """
        if not self.search_index_exists(index_name):
            self.create_vector_search_index(collection, index_name)

    def get_collection(self, collection_name: Optional[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:
            The collection object (Collection)
        """
        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:
            self.active_collection = self.scope.collection(collection_name)

        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.
        """
        try:
            collection_mgr = self.bucket.collections()
            collection_mgr.drop_collection(self.scope.name, collection_name)
        except Exception as e:
            logger.error(f"Error deleting collection: {e}")

    def create_vector_search_index(
        self,
        collection,
        index_name: Optional[str] = "vector_index",
        similarity: Literal["l2_norm", "dot_product"] = "dot_product",
    ) -> None:
        """Create a vector search index in the collection."""
        search_index_mgr = self.scope.search_indexes()
        dims = self._get_embedding_size()
        index_definition = {
            "type": "fulltext-index",
            "name": index_name,
            "sourceType": "couchbase",
            "sourceName": self.bucket.name,
            "planParams": {"maxPartitionsPerPIndex": 1024, "indexPartitions": 1},
            "params": {
                "doc_config": {
                    "docid_prefix_delim": "",
                    "docid_regexp": "",
                    "mode": "scope.collection.type_field",
                    "type_field": "type",
                },
                "mapping": {
                    "analysis": {},
                    "default_analyzer": "standard",
                    "default_datetime_parser": "dateTimeOptional",
                    "default_field": "_all",
                    "default_mapping": {"dynamic": True, "enabled": False},
                    "default_type": "_default",
                    "docvalues_dynamic": False,
                    "index_dynamic": True,
                    "store_dynamic": True,
                    "type_field": "_type",
                    "types": {
                        f"{self.scope.name}.{collection.name}": {
                            "dynamic": False,
                            "enabled": True,
                            "properties": {
                                "embedding": {
                                    "dynamic": False,
                                    "enabled": True,
                                    "fields": [
                                        {
                                            "dims": dims,
                                            "index": True,
                                            "name": "embedding",
                                            "similarity": similarity,
                                            "type": "vector",
                                            "vector_index_optimized_for": "recall",
                                        }
                                    ],
                                },
                                "metadata": {"dynamic": True, "enabled": True},
                                "content": {
                                    "dynamic": False,
                                    "enabled": True,
                                    "fields": [
                                        {
                                            "include_in_all": True,
                                            "index": True,
                                            "name": "content",
                                            "store": True,
                                            "type": "text",
                                        }
                                    ],
                                },
                            },
                        }
                    },
                },
                "store": {"indexType": "scorch", "segmentVersion": 16},
            },
            "sourceParams": {},
        }

        search_index_def = SearchIndex.from_json(json.dumps(index_definition))
        max_attempts = 10
        attempt = 0
        while attempt < max_attempts:
            try:
                search_index_mgr.upsert_index(search_index_def)
                break
            except Exception as e:
                logger.debug(f"Attempt {attempt + 1}/{max_attempts}: Error creating search index: {e}")
                time.sleep(3)
                attempt += 1

        if attempt == max_attempts:
            logger.error(f"Error creating search index after {max_attempts} attempts.")
            raise RuntimeError(f"Error creating search index after {max_attempts} attempts.")

        logger.info(f"Search index {index_name} created successfully.")

    def upsert_docs(
        self, docs: list[Document], collection: "Collection", batch_size: int = DEFAULT_BATCH_SIZE, **kwargs: Any
    ) -> None:
        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.")

        for i in range(0, len(docs), batch_size):
            batch = docs[i : i + batch_size]
            docs_to_upsert = dict()
            for doc in batch:
                doc_id = doc["id"]
                embedding = self.embedding_function([
                    doc["content"]
                ]).tolist()  # Gets new embedding even in case of document update
                doc_content = {
                    TEXT_KEY: doc["content"],
                    "metadata": doc.get("metadata", {}),
                    EMBEDDING_KEY: embedding,
                    "id": doc_id,
                }
                docs_to_upsert[doc_id] = doc_content
            collection.upsert_multi(docs_to_upsert)

    def insert_docs(
        self,
        docs: list[Document],
        collection_name: str = None,
        upsert: bool = False,
        batch_size: int = DEFAULT_BATCH_SIZE,
        **kwargs: Any,
    ) -> None:
        """Insert Documents and Vector Embeddings into the collection of the vector database. Documents are upserted in all cases."""
        if not docs:
            logger.info("No documents to insert.")
            return

        collection = self.get_collection(collection_name)
        self.upsert_docs(docs, collection, batch_size=batch_size)

    def update_docs(
        self, docs: list[Document], collection_name: str = None, batch_size: int = DEFAULT_BATCH_SIZE, **kwargs: Any
    ) -> None:
        """Update documents, including their embeddings, in the Collection."""
        collection = self.get_collection(collection_name)
        self.upsert_docs(docs, collection, batch_size)

    def delete_docs(
        self, ids: list[ItemID], collection_name: str = None, batch_size: int = DEFAULT_BATCH_SIZE, **kwargs
    ):
        """Delete documents from the collection of the vector database."""
        collection = self.get_collection(collection_name)
        # based on batch size, delete the documents
        for i in range(0, len(ids), batch_size):
            batch = ids[i : i + batch_size]
            collection.remove_multi(batch)

    def get_docs_by_ids(
        self,
        ids: Optional[list[ItemID]] = None,
        collection_name: str = None,
        include: Optional[list[str]] = None,
        **kwargs: Any,
    ) -> list[Document]:
        """Retrieve documents from the collection of the vector database based on the ids."""
        if include is None:
            include = [TEXT_KEY, "metadata", "id"]
        elif "id" not in include:
            include.append("id")

        collection = self.get_collection(collection_name)
        if ids is not None:
            docs = [collection.get(doc_id) for doc_id in ids]
        else:
            # Get all documents using couchbase query
            include_str = ", ".join(include)
            query = f"SELECT {include_str} FROM {self.bucket.name}.{self.scope.name}.{collection.name}"
            result = self.cluster.query(query)
            docs = []
            for row in result:
                docs.append(row)

        return [{k: v for k, v in doc.items() if k in include or k == "id"} for doc in docs]

    def retrieve_docs(
        self,
        queries: list[str],
        collection_name: str = None,
        n_results: int = 10,
        distance_threshold: float = -1,
        **kwargs: Any,
    ) -> QueryResults:
        """Retrieve documents from the collection of the vector database based on the queries.
        Note: Distance threshold is not supported in Couchbase FTS.
        """

        results: QueryResults = []
        for query_text in queries:
            query_vector = np.array(self.embedding_function([query_text])).tolist()[0]
            query_result = self._vector_search(
                query_vector,
                n_results,
                **kwargs,
            )
            results.append(query_result)
        return results

    def _vector_search(
        self, embedding_vector: list[float], n_results: int = 10, **kwargs
    ) -> list[tuple[dict[str, Any], float]]:
        """Core vector search using Couchbase FTS."""

        search_req = search.SearchRequest.create(
            VectorSearch.from_vector_query(
                VectorQuery(
                    EMBEDDING_KEY,
                    embedding_vector,
                    n_results,
                )
            )
        )

        search_options = SearchOptions(limit=n_results, fields=["*"])
        result = self.scope.search(self.index_name, search_req, search_options)

        docs_with_score = []

        for row in result.rows():
            doc = row.fields
            doc["id"] = row.id
            score = row.score

            docs_with_score.append((doc, score))

        return docs_with_score