import os
from langchain_community.vectorstores import FAISS
from langchain.embeddings import OpenAIEmbeddings


def embed(input_strings):
    vectorstore = FAISS.from_texts(texts=input_strings, embedding=OpenAIEmbeddings())
    return vectorstore


# Function to save a FAISS vectorstore to a specified path
def save_local(vectorstore, path="safe/"):
    if not os.path.exists(path):
        os.makedirs(path)
    file_path = os.path.join(path, "faiss_index.index")
    vectorstore.save_local(file_path)
    print(f"FAISS vectorstore saved to {file_path}")


# Function to load a FAISS vectorstore from a specified path
def load_vectorstore(path):
    embeddings = OpenAIEmbeddings()  # Needed to initialize the FAISS properly
    vectorstore = FAISS.load_local(path, embeddings, allow_dangerous_deserialization=True)
    print(f"FAISS vectorstore loaded from {path}")
    return vectorstore


# Example usage
if __name__ == "__main__":
    # Embed a few words
    words = ["hello", "world", "sample", "text"]
    faiss_db1 = embed(words)

    # Save the vectorstore
    save_local(faiss_db1)

    # Load the vectorstore
    loaded_db1 = load_vectorstore("safe/faiss_index.index")

    # Embed another set of words and create a second vectorstore
    more_words = ["FAISS", "database", "information", "retrieval"]
    faiss_db2 = embed(more_words)

    loaded_db1.merge_from(faiss_db2)
    print("Merged vectorstore with other vectorstore containing total vectors:", loaded_db1.index.ntotal)