from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_community.llms import OpenAI
from langchain_community.chat_models import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain.memory.buffer import ConversationBufferMemory
import os
import streamlit as st
from dotenv import load_dotenv
from langchain.chains import RetrievalQAWithSourcesChain, ConversationalRetrievalChain
from langchain_community.llms import HuggingFaceEndpoint


class LangChainTools:
    """
    Esta clase maneja algunas herramientas integraciones con las que
    cuenta LangChain.
    """

    def load_embedding_opnai(self):
        """Esta funcion carga un modelo de embedding de OpenAI

        Returns:
            _type_: Retorno a un objetito de tipo embedding de OpenAI
        """

        # Cargamos la variable que contiene la api_key de OpenAI
        load_dotenv()
        openai_api_key = os.getenv("api_key")
        # Define an OpenAI embeddings model
        self.embedding_model = OpenAIEmbeddings(openai_api_key=openai_api_key)
        # st.success('El modelo de embeddins de OpneAI se ha cargado')

        return self.embedding_model

    @st.cache_resource
    def create_vector_strore(
        _self, _docs_split: list, _file_name: str, _embedding_model
    ):
        """Esta funcion construye un vector store a partir de un documento

        Args:
            _docs_split (list): Lista de documentos divididos
            _file_name (str): Nombre del documento
        """

        db_name = _file_name.replace(".pdf", "").replace(" ", "_").lower()

        # Cargamos el modelo de embeddings
        # _embedding_model = self._embedding_model

        # Verificamos si existe la vector strore
        persist_directory = f"embeddings/{db_name}"

        if os.path.exists(persist_directory):
            vectordb = Chroma(
                persist_directory=persist_directory, embedding_function=_embedding_model
            )
        else:
            vectordb = Chroma.from_documents(
                persist_directory=persist_directory,
                documents=_docs_split,
                embedding=_embedding_model,
            )

            vectordb.persist()

        return vectordb

    def load_llm_openai(self):
        """Esta funcion carga un modelo de LLM de OpenAI

        Returns:
            _type_: Retorno a un objetito de tipo LLM de OpenAI
        """

        # Cargamos la variable que contiene la api_key de OpenAI
        load_dotenv()
        openai_api_key = os.getenv("api_key")

        temperature = 0.5
        llm_openai = ChatOpenAI(
            model_name="gpt-3.5-turbo",
            temperature=temperature,
            openai_api_key=openai_api_key,
            max_tokens=1000,
        )

        return llm_openai

    def load_llm_open_source(self):
        """Esta funcion carga un modelo de LLM OpenSource desde HuggingFace

        Returns:
            _type_: Retorno a un objetito de tipo LLM de OpenAI
        """
        # model_huggingface = "google/gemma-1.1-7b-it"  # Es buena y funciona en espanol
        # model_huggingface = (
        #    "google/gemma-1.1-2b-it"  # Es buena y funciona en espanol funciona rapido
        # )
        # model_huggingface = 'tiiuae/falcon-7b-instruct'
        # model_huggingface = 'mistralai/Mistral-7B-Instruct-v0.2'
        huggingfacehub_api_token = "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"

        model_huggingface = "mistralai/Mixtral-8x7B-Instruct-v0.1"  # Es buena y funciona en espanol funciona rapido

        # Define the LLM
        llm = HuggingFaceEndpoint(
            repo_id=model_huggingface,
            huggingfacehub_api_token=huggingfacehub_api_token,
            temperature=0.5,
            max_new_tokens=500,
        )

        return llm

    def load_prompt_template(self):
        """Esta funcion construye un prompt template de lanfchain.

        Returns:
            _type_: Retorno a un prompt template de LangChain.
        """
        template = """Responde a la siguiente pregunta utilizando los documentos proporcionados y citando las fuentes relevantes entre corchetes []:

        Pregunta: {question}

        Respuesta:"""

        prompt_template = PromptTemplate(
            template=template, input_variables=["question"]
        )

        return prompt_template

    def define_retrieval_qa(self, _llm, _vectordb, _file_name, _embedding_model):
        """Esta función integra un LLM y una base de datos vectorial en una
        chain de LangChain para hacer requerimientos. Este modelo no integra memoria.

        Args:
            _llm (_type_): <Modelo Largo de Lenguaje.
            _vectordb (_type_): Base de datos vectorial.
            _file_name (_type_): Nombre del archvio con la que se crea la BDV
            para cargarla si existe.
            _embedding_model (_type_): Modelo de embedding.

        Returns:
            _type_: Retorna un objeto RetrievalQAWithSourcesChain con el quie
            se pueden hacer requerimientos a la chain que integra el modelo
            y la BDV.
        """

        db_name = _file_name.replace(".pdf", "").replace(" ", "_").lower()

        # Verificamos si existe la vector strore
        persist_directory = f"embeddings/{db_name}"

        _vectordb = Chroma(
            persist_directory=persist_directory, embedding_function=_embedding_model
        )

        # Define the Retrieval QA Chain to integrate the database and LLM
        qa = RetrievalQAWithSourcesChain.from_chain_type(
            _llm,
            retriever=_vectordb.as_retriever(),
            return_source_documents=True,  # Devuelve los documentos fuente
            max_tokens_limit=1000,  # Límite máximo de tokens para el LLM
            reduce_k_below_max_tokens=True,  # Reduce k si los tokens exceden el límite
            verbose=True,  # Modo verboso
        )

        return qa

    @st.cache_resource
    def define_retrieval_qa_memory(
        _self, _llm, _vectordb, _file_name, _embedding_model
    ):
        """Esta función integra un LLM y una base de datos vectorial en una
        chain de LangChain para hacer requerimientos. Este modelo integra memoria.

        Args:
            _llm (_type_): <Modelo Largo de Lenguaje.
            _vectordb (_type_): Base de datos vectorial.
            _file_name (_type_): Nombre del archvio con la que se crea la BDV
            para cargarla si existe.
            _embedding_model (_type_): Modelo de embedding.

        Returns:
            _type_: Retorna un objeto RetrievalQAWithSourcesChain con el quie
            se pueden hacer requerimientos a la chain que integra el modelo
            y la BDV.
        """

        db_name = _file_name.replace(".pdf", "").replace(" ", "_").lower()

        # Verificamos si existe la vector strore
        persist_directory = f"embeddings/{db_name}"

        _vectordb = Chroma(
            persist_directory=persist_directory, embedding_function=_embedding_model
        )

        # Configura la memoria
        memory = ConversationBufferMemory(
            memory_key="chat_history", return_messages=True, output_key="answer"
        )

        # Define the Retrieval QA Chain to integrate the database and LLM
        conversation = ConversationalRetrievalChain.from_llm(
            _llm,
            retriever=_vectordb.as_retriever(),
            memory=memory,
            verbose=True,  # Modo verboso
            return_source_documents=True,  # Devuelve los documentos fuente
        )

        template = """Utiliza los siguientes fragmentos de contexto para responder la pregunta al final. Si no sabes la respuesta, simplemente di que no sabes, no intentes inventar una respuesta. La respuesta dala con un formateo de markdown. Responde a la pregunta siempre en español.
        
        {context}

        Pregunta: {question}
        Respuesta:"""
        conversation.combine_docs_chain.llm_chain.prompt.template = template

        return conversation