##Instalación de paquetes necesarios
import streamlit as st
import os
import time
import torch
from utils import *
from dotenv import load_dotenv
load_dotenv()

##import nest_asyncio
##nest_asyncio.apply()
from llama_parse import LlamaParse

from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import VectorStoreIndex, ServiceContext
from llama_index.core import SimpleDirectoryReader
from llama_index.core import Settings
######   

## titulos y cabeceras 
st.set_page_config('compare PDF por LLM')
st.title("Comparar PDFs mediante LLM")
st.subheader("Campos a comparar en tu PDF",divider='rainbow')

OPENAI_API_KEY = st.text_input('OpenAI API Key', type='password')
LLAMA_CLOUD_API_KEY = st.text_input('LLAMA API Key', type='password')

####  Inicializar mensajes de chat 
if "messages" not in st.session_state.keys():
    st.session_state.messages = [
        {"role": "assistant", "content": "Ask me a question about PDFs files you provided me"}
    ]

@st.cache_resource(show_spinner=False) # Añade decorador de caché
def cargar_embedmodel_y_llmmodel():
    return True
    
#esta variable es para tener aqui un listado de aquellos ficheros que se han ido subiendo
archivos = [] 

## carga y almacenamiento de ficheros almacenada, acepta varios.
with st.sidebar:
    archivos = load_name_files(FILE_LIST)
    files_uploaded = st.file_uploader(
        "Carga tus ficheros PDF",
        type="pdf",
        accept_multiple_files=True,
        on_change=st.cache_resource.clear
        )  

    if st.button("Guardar y procesar por LLM", type="secondary",help="donde buscará lo que comparará"):
        for pdf in files_uploaded:
            if pdf is not None and pdf.name not in archivos:
                archivos.append(pdf.name)
             
        archivos = save_name_files(FILE_LIST, archivos)
    

    if len(archivos)>0:
        st.write('Los archivos PDF se han cargados:')
        lista_documentos = st.empty()
        with lista_documentos.container():
            for arch in archivos:
                st.write(arch)
            if st.button('Borrar ficheros'):
                archivos = []
                clean_files(FILE_LIST)
                lista_documentos.empty()


# comprueba que hay archivos a ser tratados
if len(archivos)>0:
# comprueba que hay consulta a responder   
   if user_question := st.chat_input("Realizar consulta:"): 
      st.session_state.messages.append({"role": "user", "content": user_question})

   if user_question: 
        for message in st.session_state.messages: # Muestra anteriores mensajes
            with st.chat_message(message["role"]):
              st.write(message["content"])

        alert = st.warning("Sea paciente") # Mensaje de aviso o warning al usuario
        time.sleep(3) # establece tiempo espera en 3 segundos
        alert.empty() # borra el aviso 

# se define el analizador-parser de los documentos.
        os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
        parser = LlamaParse(
##      api_key=os.environ["LLAMA_CLOUD_API_KEY"],    ##API de acceso a Cloud de LlamaIndex
        api_key=LLAMA_CLOUD_API_KEY,     
        result_type="markdown",                       # se toma "markdown", tambien hay text disponible 
        verbose=True,
        )

        cargar_embedmodel_y_llmmodel()

#se parametrizan los modelos de embedding y LLM        
        embed_model=OpenAIEmbedding(model="text-embedding-3-small")   #embeddings para base de conocimiento
        llm = OpenAI(model="gpt-3.5-turbo-0125")                      #modelo LLM usado

        Settings.llm = llm
        Settings.embed_model = embed_model

        
        tratar = load_name_files(FILE_LIST) ##variable que tomará los ficheros a tratar recuperados de funcion
#        st.write(tratar[0]) # se puede desasteriscar en desarrollo para apoyo
#        st.write(tratar[1]) # se puede desasteriscar en desarrollo para apoyo    
        

# Carga de los ficheros mediante LlamaParse, se ejecutará job para cada analizador-parser de los mismos        
        docs_202401 = parser.load_data( f'{tratar[0]}')     
        docs_202402 = parser.load_data( f'{tratar[1]}')
        
#uso de MarkdownElementNodeParser para analizar la salida de LlamaParse mediante un motor de consultas de recuperación(recursivo)
        from llama_index.core.node_parser import MarkdownElementNodeParser
        node_parser = MarkdownElementNodeParser(llm=OpenAI(model="gpt-3.5-turbo-0125"), num_workers=8)  

        import pickle
        from llama_index.postprocessor.flag_embedding_reranker import FlagEmbeddingReranker
# se parametriza el modelo reranker
        reranker = FlagEmbeddingReranker(
            top_n=5,
            model="BAAI/bge-reranker-large",
        )
#funcion para Facilitar el motor de consultas sobre el almacén de vectores, y poderse realizar la recuperación.
        def create_query_engine_over_doc(docs, nodes_save_path=None):
            """Big function to go from document path -> recursive retriever."""
            if nodes_save_path is not None and os.path.exists(nodes_save_path):
                raw_nodes = pickle.load(open(nodes_save_path, "rb"))
            else:
                raw_nodes = node_parser.get_nodes_from_documents(docs)
                if nodes_save_path is not None:
                    pickle.dump(raw_nodes, open(nodes_save_path, "wb"))

            base_nodes, objects = node_parser.get_nodes_and_objects(
                raw_nodes
            )

### Recuperador-retriever
            # indice y motor 
            vector_index = VectorStoreIndex(nodes=base_nodes+objects)
            query_engine = vector_index.as_query_engine(
                similarity_top_k=15,
                node_postprocessors=[reranker]
            )
            return query_engine, base_nodes, vector_index   ###devuelve motor de consultas y nodos
        
## motores de consulta y nodos para cada documento usando la función anterior.
## En los ficheros .pkl se puede ver la estructura de los documentos que ha conformado o analizado y será con la que trabajará.
        query_engine_202401, nodes_202401,vindex1 = create_query_engine_over_doc(
        docs_202401, nodes_save_path="202401_nodes.pkl"
        )
        query_engine_202402, nodes_202402,vindex2 = create_query_engine_over_doc(
        docs_202402, nodes_save_path="202402_nodes.pkl"
        )

        from llama_index.core.tools import QueryEngineTool, ToolMetadata
        from llama_index.core.query_engine import SubQuestionQueryEngine

        from llama_index.core.llms import ChatMessage

# motor de consulta como tool, configuración y contexto de los datos que deberá proveer por los que será consultado
# debajo se usa como motor de subconsultas SubQuestionQueryEngine
        query_engine_tools = [
            QueryEngineTool(
                query_engine=query_engine_202401,
                metadata=ToolMetadata(
                    name="pdf_ENERO",
                    description=(
#                        "Provides information about Datos del Producto for ENERO"
#                        "Provides information about values of fields of Datos del Producto, Titular, Fabricante,Composicion,Envases, Usos y Dosis Autorizados,Plazos de Seguridad"
                                                  """\
                            The documents provided are plant protection product data sheets in PDF format.
                            Provides information about values of fields of Datos del Producto, Titular, Fabricante,Composicion,Envases, 
                            Usos y Dosis Autorizados,Plazos de Seguridad:
                            # Datos del Producto
                            |Numero de Registro|
                            |Estado|
                            |Fechas Inscripción|
                            |Renovación|
                            |Caducidad|
                            |Nombre Comercial|

                            # Titular

                            # Fabricante

                            # Composición

                            # Envases

                            # Usos y Dosis Autorizados
                            |USO|
                            |AGENTE|
                            |Dosis|
                            |Condic. Especifico|

                            """
                    ),
                ),
            ),
            QueryEngineTool(
                query_engine=query_engine_202402,
                metadata=ToolMetadata(
                    name="pdf_FEBRERO",
                    description=(
#                        "Provides information about Datos del Producto for FEBRERO"
#                        "Provides information about values of fields of Datos del Producto, Titular, Fabricante,Composicion,Envases, Usos y Dosis Autorizados,Plazos de Seguridad"
                         """\
                            The documents provided are plant protection product data sheets in PDF format.
                            Provides information about values of fields of Datos del Producto, Titular, Fabricante,Composicion,Envases, 
                            Usos y Dosis Autorizados,Plazos de Seguridad:
                            # Datos del Producto
                            |Numero de Registro|
                            |Estado|
                            |Fechas Inscripción|
                            |Renovación|
                            |Caducidad|
                            |Nombre Comercial|

                            # Titular

                            # Fabricante

                            # Composición

                            # Envases

                            # Usos y Dosis Autorizados
                            |USO|
                            |AGENTE|
                            |Dosis|
                            |Condic. Especifico|

                            """
                    ),
                ),
            ),
        ]
# subconsultas con tool creada a través de SubQuestionQueryEngine
        sub_query_engine = SubQuestionQueryEngine.from_defaults(
            query_engine_tools=query_engine_tools,
            llm=llm
        )

        if "chat_engine" not in st.session_state.keys(): # Initializa motor chat
# para que generen las subconsultas con la consulta-query del usuario           
            streaming_response = sub_query_engine.query(user_question)
        
##      If last message is not from assistant, generate a new response
        if st.session_state.messages[-1]["role"] != "assistant":
            with st.chat_message("assistant"):
                with st.spinner("Thinking..."):    #figura del spinner de streamlit mientras se ejecuta bloque
                     response = st.write(streaming_response.response)  #respuesta entregada a la query-consulta del usuario
                     st.session_state.messages.append({"role": "assistant", "content": response})