# gradio_interface.py
import gradio as gr
from rag_functions import *

def demo():
    with gr.Blocks(theme="base") as demo:
        vector_db = gr.State()
        qa_chain = gr.State()
        collection_name = gr.State()
        
        gr.Markdown(
        """<center><h2>Assistente de Análise de Relatórios de Metrologia</center></h2>
        <h3>Faça qualquer pergunta sobre seus relatórios de metrologia</h3>""")
        gr.Markdown(
        """<b>Nota:</b> Este assistente de IA, utilizando Langchain e LLMs de código aberto, realiza geração aumentada por recuperação (RAG) a partir de seus relatórios de metrologia em formato PDF. \
        A interface do usuário está organizada para facilitar o entendimento do fluxo de trabalho do RAG. Este chatbot leva em consideração perguntas anteriores ao gerar respostas, e inclui referências documentais para maior clareza.<br>
        <br><b>Aviso:</b> Este espaço usa a CPU básica gratuita do Hugging Face. Algumas etapas e modelos LLM utilizados abaixo (pontos finais de inferência gratuitos) podem levar algum tempo para gerar uma resposta.
        """)
        
        with gr.Tab("Etapa 1 - Carregar Relatórios"):
            with gr.Row():
                document = gr.Files(height=100, file_count="multiple", file_types=["pdf"], interactive=True, label="Carregue seus relatórios de metrologia (único ou múltiplos)")
        
        with gr.Tab("Etapa 2 - Processar Relatórios"):
            with gr.Row():
                db_btn = gr.Radio(["ChromaDB"], label="Tipo de banco de dados vetorial", value = "ChromaDB", type="index", info="Escolha o banco de dados vetorial")
            with gr.Accordion("Opções avançadas - Divisor de texto do relatório", open=False):
                with gr.Row():
                    slider_chunk_size = gr.Slider(minimum = 100, maximum = 1000, value=600, step=20, label="Tamanho do bloco", info="Tamanho do bloco", interactive=True)
                with gr.Row():
                    slider_chunk_overlap = gr.Slider(minimum = 10, maximum = 200, value=40, step=10, label="Sobreposição do bloco", info="Sobreposição do bloco", interactive=True)
            with gr.Row():
                db_progress = gr.Textbox(label="Inicialização do banco de dados vetorial", value="Nenhum")
            with gr.Row():
                db_btn = gr.Button("Gerar banco de dados vetorial")
            
        with gr.Tab("Etapa 3 - Inicializar cadeia de QA"):
            with gr.Row():
                llm_btn = gr.Radio(list_llm_simple, \
                    label="Modelos LLM", value = list_llm_simple[0], type="index", info="Escolha seu modelo LLM")
            with gr.Accordion("Opções avançadas - Modelo LLM", open=False):
                with gr.Row():
                    slider_temperature = gr.Slider(minimum = 0.01, maximum = 1.0, value=0.7, step=0.1, label="Temperatura", info="Temperatura do modelo", interactive=True)
                with gr.Row():
                    slider_maxtokens = gr.Slider(minimum = 224, maximum = 4096, value=1024, step=32, label="Máximo de Tokens", info="Máximo de tokens do modelo", interactive=True)
                with gr.Row():
                    slider_topk = gr.Slider(minimum = 1, maximum = 10, value=3, step=1, label="Amostras top-k", info="Amostras top-k do modelo", interactive=True)
            with gr.Row():
                llm_progress = gr.Textbox(value="Nenhum",label="Inicialização da cadeia QA")
            with gr.Row():
                qachain_btn = gr.Button("Inicializar cadeia de Pergunta e Resposta")

        with gr.Tab("Etapa 4 - Chatbot"):
            chatbot = gr.Chatbot(height=300)
            with gr.Accordion("Avançado - Referências do Relatório", open=False):
                with gr.Row():
                    doc_source1 = gr.Textbox(label="Referência 1", lines=2, container=True, scale=20)
                    source1_page = gr.Number(label="Página", scale=1)
                with gr.Row():
                    doc_source2 = gr.Textbox(label="Referência 2", lines=2, container=True, scale=20)
                    source2_page = gr.Number(label="Página", scale=1)
                with gr.Row():
                    doc_source3 = gr.Textbox(label="Referência 3", lines=2, container=True, scale=20)
                    source3_page = gr.Number(label="Página", scale=1)
            with gr.Row():
                msg = gr.Textbox(placeholder="Digite a mensagem (exemplo: 'Qual a precisão dos instrumentos utilizados?')", container=True)
            with gr.Row():
                submit_btn = gr.Button("Enviar mensagem")
                clear_btn = gr.ClearButton([msg, chatbot], value="Limpar conversa")
            
        # Eventos de pré-processamento
        db_btn.click(initialize_database, \
            inputs=[document, slider_chunk_size, slider_chunk_overlap], \
            outputs=[vector_db, collection_name, db_progress])
        qachain_btn.click(initialize_LLM, \
            inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], \
            outputs=[qa_chain, llm_progress]).then(lambda:[None,"",0,"",0,"",0], \
            inputs=None, \
            outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
            queue=False)

        # Eventos do Chatbot
        msg.submit(conversation, \
            inputs=[qa_chain, msg, chatbot], \
            outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
            queue=False)
        submit_btn.click(conversation, \
            inputs=[qa_chain, msg, chatbot], \
            outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
            queue=False)
        clear_btn.click(lambda:[None,"",0,"",0,"",0], \
            inputs=None, \
            outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
            queue=False)
    demo.queue().launch(debug=True)


if __name__ == "__main__":
    demo()