import logging
import os
import requests



from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings


class RAG:
    NO_ANSWER_MESSAGE: str = "Ho sento, no he pogut respondre la teva pregunta."

    #vectorstore = "index-intfloat_multilingual-e5-small-500-100-CA-ES" # mixed
    #vectorstore = "vectorestore" # CA only
    vectorstore = "index-BAAI_bge-m3-1500-200-recursive_splitter-CA_ES_UE"

    def __init__(self, hf_token, embeddings_model, model_name):

        self.model_name = model_name
        self.hf_token = hf_token
        
        # load vectore store
        embeddings = HuggingFaceEmbeddings(model_name=embeddings_model, model_kwargs={'device': 'cpu'})
        self.vectore_store = FAISS.load_local(self.vectorstore, embeddings, allow_dangerous_deserialization=True)#, allow_dangerous_deserialization=True)

        logging.info("RAG loaded!")
    
    def get_context(self, instruction, number_of_contexts=2):

        documentos = self.vectore_store.similarity_search_with_score(instruction, k=number_of_contexts)

        return documentos
        
    def predict(self, instruction, sys_prompt, context, model_parameters):

        from openai import OpenAI
        
        # init the client but point it to TGI
        client = OpenAI(
            base_url=os.getenv("MODEL")+ "/v1/",
            api_key=os.getenv("HF_TOKEN")
        )

        #sys_prompt = "You are a helpful assistant. Answer the question using only the context you are provided with. If it is not possible to do it with the context, just say 'I can't answer'. <|endoftext|>"
        #query = f"Context:\n{context}\n\nQuestion:\n{instruction}"
        query = f"Context:\n{context}\n\nQuestion:\n{instruction}\n\n{sys_prompt}"
        print(query)
        #query = f"{sys_prompt}\n\nQuestion:\n{instruction}\n\nContext:\n{context}"
        chat_completion = client.chat.completions.create(
            model="tgi",
            messages=[
                #{"role": "system", "content": sys_prompt },
                {"role": "user", "content": query}
            ],
            max_tokens=model_parameters['max_new_tokens'], # TODO: map other parameters
            frequency_penalty=model_parameters['repetition_penalty'], # this doesn't appear to do much, not a replacement for repetition penalty
            # presence_penalty=model_parameters['repetition_penalty'],
            # extra_body=model_parameters,
            stream=False,
            stop=["<|im_end|>", "<|end_header_id|>", "<|eot_id|>", "<|reserved_special_token"]
        )
        return(chat_completion.choices[0].message.content)

    
    def beautiful_context(self, docs):

        text_context = ""

        full_context = ""
        source_context = []
        for doc in docs:
            text_context += doc[0].page_content
            full_context += doc[0].page_content + "\n"
            full_context += doc[0].metadata["Títol de la norma"] + "\n\n"
            full_context += doc[0].metadata["url"] + "\n\n"
            source_context.append(doc[0].metadata["url"])

        return text_context, full_context, source_context

    def get_response(self, prompt: str, sys_prompt: str, model_parameters: dict) -> str:
        try:
            docs = self.get_context(prompt, model_parameters["NUM_CHUNKS"])
            text_context, full_context, source = self.beautiful_context(docs)

            del model_parameters["NUM_CHUNKS"]

            response = self.predict(prompt, sys_prompt, text_context, model_parameters)

            if not response:
                return self.NO_ANSWER_MESSAGE

            return response, full_context, source
        except Exception as err:
            print(err)