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.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
from langchain_community.embeddings import HuggingFaceEmbeddings
from streamlit.runtime.state import session_state
from langchain_community.llms import Ollama


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

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

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

        huggingfacehub_api_token = "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"

        model_name = "sentence-transformers/all-mpnet-base-v2"
        model_kwargs = {"device": "cpu"}
        encode_kwargs = {"normalize_embeddings": False}

        self.embedding_model = HuggingFaceEmbeddings(
            model_name=model_name,
            model_kwargs=model_kwargs,
            encode_kwargs=encode_kwargs,
        )

        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()
        # st.success(_file_name)

        if "db_name" not in st.session_state.keys():
            st.session_state.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 "persist_director" not in st.session_state.keys():
            st.session_state.persist_directory = (
                f"embeddings/{st.session_state.db_name}"
            )

        if os.path.exists(st.session_state.persist_directory):
            vectordb = Chroma(
                persist_directory=st.session_state.persist_directory,
                embedding_function=_embedding_model,
            )
        else:
            vectordb = Chroma.from_documents(
                persist_directory=st.session_state.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

    @st.cache_resource
    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"
        # model_huggingface = 'mistralai/Mixtral-8x7B-Instruct-v0.1'
        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.1,
            max_new_tokens=1000,
        )

        return llm

    @st.cache_resource
    def load_llm_ollama(_self):
        """Esta funcion carga un modelo de LLM OpenSource desde Ollama

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

        # Elegimos el modelo de Ollama que utilizaremos

        model: str = "gemma:2b"

        llm = Ollama(
            model=model,
            temperature=0.1,
            num_ctx=1000,
        )

        return llm

    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, _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=st.session_state.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 en español la pregunta al final. Si no sabes la respuesta, simplemente di que no sabes, no intentes inventar una respuesta.

        {context}

        Pregunta: {question}
        Respuesta:"""

        # template = """Utiliza los siguientes fragmentos de contexto como ejemplo para responder la pregunta al final. Organiza tu respuesta de manera clara y concisa, proporcionando información relevante y evitando divagaciones innecesarias.

        # {context}

        # Pregunta: {question}
        # Respuesta en español:"""

        # conversation.combine_docs_chain.llm_chain.prompt.template = template
        # conversation.question_generator.prompt.template = "Dado el siguiente diálogo y una pregunta de seguimiento, reformula la pregunta de seguimiento para que sea una pregunta independiente, en su idioma original.\n\nHistorial del chat:\n{chat_history}\nPregunta de seguimiento: {question}\nPregunta independiente:"

        return conversation