from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
# from langchain_huggingface import HuggingFaceEmbeddings
from huggingface_hub import snapshot_download
import logging


class VectorStore:
    def __init__(self, embeddings_model, vs_local_path=None, vs_hf_path=None, number_of_contexts=2):
        
        self.number_of_contexts = number_of_contexts

        logging.info("Loading vectorstore...")        
        
        embeddings = HuggingFaceEmbeddings(model_name=embeddings_model)
        logging.info(f"Loaded embeddings model: {embeddings_model}")


        if vs_hf_path:
            hf_vectorstore = snapshot_download(repo_id=vs_hf_path)
            self.vectore_store = FAISS.load_local(hf_vectorstore, embeddings, allow_dangerous_deserialization=True)
            logging.info(f"Loaded vectorstore from {vs_hf_path}")
        else:
            self.vectore_store = FAISS.load_local(vs_local_path, embeddings, allow_dangerous_deserialization=True)
            logging.info(f"Loaded vectorstore from {vs_local_path}")

    
    def get_context(self, instruction, number_of_contexts=2):

        logging.info(f"Getting context for instruction: {instruction}")
        documentos = self.vectore_store.similarity_search_with_score(instruction, k=self.number_of_contexts)
        return self._beautiful_context(documentos)
    
    def _beautiful_context(self, docs):
        context = ""
        for doc in docs:
            context += doc[0].page_content + "\n\n"

        print("Context: ", context)
        return context[:-1]