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)


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_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=False,  # Modo verboso
            return_source_documents=True  # Devuelve los documentos fuente
        )

        return conversation