class Document(BaseModel):
    """
    A document is a chunk of text.

    content (str): the string that makes up the chunk of text
    id (str): the optional id associated with the chunk of text
    metadata (Optional[dict]): optional information about the chunk of text
    """
    content: str
    id: Optional[str] = Field(default_factory=lambda: str(uuid.uuid4()))
    metadata: Optional[dict] = Field(default_factory=dict)

class LangChainWiseAgentVectorDB(WiseAgentVectorDB):
    """
    An abstract class that makes use of a LangChain vector database.
    """

    def __init__(self, embedding_model_name: Optional[str] = DEFAULT_EMBEDDING_MODEL_NAME):
        """
        Initialize a new instance of LangChainWiseAgentVectorDB.


        Args:
            embedding_model_name (Optional[str]): the optional name of the embedding model to use
        """
        self._embedding_model_name = embedding_model_name
        self._embedding_function = HuggingFaceEmbeddings(model_name=self.embedding_model_name)

    @property
    def embedding_model_name(self):
        """Get the name of the embedding model."""
        return self._embedding_model_name

    def get_embedding_function(self):
        """Get the embedding function."""
        if not hasattr(self, "_embedding_function"):
            # instances populated from PyYAML won't have this set initially
            self._embedding_function = HuggingFaceEmbeddings(model_name=self.embedding_model_name)
        return self._embedding_function

    def convert_from_lang_chain_documents(self, documents: List[LangChainDocument]) -> List[Document]:
        return [Document(content=document.page_content, metadata=document.metadata) for document in documents]

    @abstractmethod
    def get_or_create_collection(self, collection_name: str):
        ...

    @abstractmethod
    def delete_collection(self, collection_name: str):
        ...

    @abstractmethod
    def insert_documents(self, documents: List[Document], collection_name: str):
        ...

    @abstractmethod
    def insert_or_update_documents(self, documents: List[Document], collection_name: str):
        ...

    @abstractmethod
    def delete_documents(self, ids: List[str], collection_name: str):
        ...

    @abstractmethod
    def query(self, queries: List[str], collection_name: str, k: Optional[int] = 4) -> List[List[Document]]:
        ...

class PGVectorLangChainWiseAgentVectorDB(LangChainWiseAgentVectorDB):
    """
    A LangChainWiseAgentVectorDB implementation that makes use of a LangChain PGVector database.
    """

    yaml_tag = u'!PGVectorLangChainWiseAgentVectorDB'

    def __init__(self, connection_string: str, embedding_model_name: Optional[str] = DEFAULT_EMBEDDING_MODEL_NAME):
        """
        Initialize a new instance of PGVectorLangChainWiseAgentVectorDB.


        Args:
            connection_string (str): the connection string for the PGVector database
            embedding_model_name (Optional[str]): the optional name of the embedding model to use
        """
        super().__init__(embedding_model_name)
        self._connection_string = connection_string
        self._vector_dbs = {}

    def __repr__(self):
        """Return a string representation of the vector DB."""
        return (f"{self.__class__.__name__}(connection_string={self.connection_string},"
                f"embedding_model_name={self.embedding_model_name})")


    def __getstate__(self) -> object:
        """Return the state of the vector DB. Removing _vector_dbs and _embedding_function to avoid them being serialized/deserialized by pyyaml."""
        state = self.__dict__.copy()
        del state['_vector_dbs']
        del state['_embedding_function']
        return state

    @property
    def connection_string(self):
        """Get the connection string."""
        return self._connection_string

    def get_or_create_collection(self, collection_name: str):
        if not hasattr(self, "_vector_dbs"):
            # instances populated from PyYAML won't have this set initially
            self._vector_dbs = {}
        if collection_name not in self._vector_dbs:
            self._vector_dbs[collection_name] = PGVector(embeddings=self.get_embedding_function(),
                                                         collection_name=collection_name,
                                                         connection=self._connection_string)

    def delete_collection(self, collection_name: str):
        self.get_or_create_collection(collection_name)
        if collection_name in self._vector_dbs:
            self._vector_dbs[collection_name].delete_collection()
            del self._vector_dbs[collection_name]

    def insert_documents(self, documents: List[Document], collection_name: str):
        self.get_or_create_collection(collection_name)
        self._vector_dbs[collection_name].add_texts(texts=[doc.content for doc in documents],
                                                    ids=[doc.id for doc in documents],
                                                    metadatas=[doc.metadata for doc in documents])

    def insert_or_update_documents(self, documents: List[Document], collection_name: str):
        self.get_or_create_collection(collection_name)
        self.insert_documents(documents, collection_name)

    def delete_documents(self, document_ids: List[str], collection_name: str):
        self.get_or_create_collection(collection_name)
        if collection_name in self._vector_dbs:
            self._vector_dbs[collection_name].delete(ids=document_ids)

    def query(self, queries: List[str], collection_name: str, k: Optional[int] = 4) -> List[List[Document]]:
        self.get_or_create_collection(collection_name)
        if collection_name in self._vector_dbs:
            return [self.convert_from_lang_chain_documents(self._vector_dbs[collection_name].similarity_search(query, k))
                    for query in queries]

class WiseAgentVectorDB(yaml.YAMLObject):
    """Abstract class to define the interface for a WiseAgentVectorDB."""
    yaml_tag = u'!WiseAgentVectorDB'

    @abstractmethod
    def get_or_create_collection(self, collection_name: str):
        """
        Get the collection for the vector DB or create it if it doesn't already exist.


        Args:
            collection_name (str): the name of the collection
        """
        ...

    @abstractmethod
    def delete_collection(self, collection_name: str):
        """
        Delete the collection with the specified name from the vector DB.


        Args:
            collection_name (str): the name of the collection to delete
        """
        ...

    @abstractmethod
    def insert_documents(self, documents: List[Document], collection_name: str):
        """
        Insert the given documents into the specified collection in the vector DB.


        Args:
            documents (List[Document]): the documents to be inserted into the specified collection
            collection_name (str): the name of the collection in the vector DB to insert the documents into
        """
        ...

    @abstractmethod
    def insert_or_update_documents(self, documents: List[Document], collection_name: str):
        """
        Insert the given documents into the specified collection in the vector DB, updating any
        documents that already exist in the collection.


        Args:
            documents (List[Document]): the documents to be inserted into the specified collection
            collection_name (str): the name of the collection in the vector DB to insert the documents into
        """
        ...

    @abstractmethod
    def delete_documents(self, ids: List[str], collection_name: str):
        """
        Delete documents from the specified collection in the vector DB.


        Args:
            ids (List[str]): the list of document IDs to be deleted
            collection_name (str): the name of the collection in the vector DB to delete the documents from
        """
        ...

    @abstractmethod
    def query(self, queries: List[str], collection_name: str, k: Optional[int]) -> List[List[Document]]:
        """
        Retrieve documents from the specified collection in the vector DB using the given queries.


        Args:
            queries (List[str]): the list of queries where each query is a string
            collection_name (str): the name of the collection in the vector DB to query
            k (Optional[int]): the number of documents to retrieve for each query

        Returns:
            List[List[Document]]: the list containing a list of documents that were
            retrieved for each query
        """
        ...