# Application file for Gradio App for OpenAI Model

import gradio as gr
import time
import datetime
import os

from lc_base.chain import openai_chain
from lc_base.dnd_database import create_dnd_database

############################# Global Params #############################

time_diff = 0
# model_name="gpt-3.5-turbo-1106" # FOR TESTING
# model_name = "gpt-4-1106-preview"
model_name = "gpt-4o-mini-2024-07-18"
search_type = "stuff"
input_question = ""
model_response = ""
user_feedback = ""

dir = ""
title = """<h1 align="center">ResearchBuddy</h1>"""
description = """<br><br><h3 align="center">This is a GPT based Research Buddy to assist in navigating new research topics.</h3>"""

DEFAULT_STATUS = "⬆Upload PDF files"


############################# Drag and Drop PDF processing #############################
def check_pdfs(pdf_files):
    global db
    db = create_dnd_database(pdf_files)
    if not db:
        return "Please upload PDF files again or submit a drive link containing only PDFs."
    else:
        return "Processing Completed - You can start the chat now!"

############################# Chatbot Specific functions #############################
def user(user_message, history):
    return "", history + [[user_message, None]]

def respond(message, chat_history):

    global time_diff, model_response, input_question

    question = str(message)
    chain = openai_chain(inp_dir=dir)

    query = question
    start_time = time.time()

    output = chain.get_response_from_drive(query=query, database=db, k=10, model_name=model_name, type=search_type)
    
    # Update global variables for logging
    time_diff = time.time() - start_time
    model_response = output
    input_question = question
    bot_message = output
    chat_history.append((message, bot_message))

    time.sleep(1) # Pause for a second to avoid overloading
    return " ", chat_history 


############################# Gradio Application Block #############################
with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald", neutral_hue="slate")) as chat:
    gr.HTML(title)

    global db

    # PDF Drag and Drop + Drive link Input + Status containers
    with gr.Row(equal_height=True):
        with gr.Column():
            with gr.Row():
                pdf_files_dnd = gr.File(file_count='multiple', height=250, label="Upload PDF Files")
            with gr.Row():
                status_message = gr.Text(label="Status", value=DEFAULT_STATUS, text_align='center')
            


    pdf_files_dnd.change(
        fn=check_pdfs, 
        inputs=[pdf_files_dnd], 
        outputs=[status_message], 
        preprocess=False, 
        postprocess=False) # Set preprocess and postprocess to False, to avoid the tmpfile object creation, instead get a Dict

    # Chatbot container
    chatbot = gr.Chatbot(height=750)
    msg = gr.Textbox(label="Send a message", placeholder="Send a message",
                             show_label=False, container=False)  
    
    with gr.Row():
        with gr.Column():
            clear_history_button = gr.ClearButton(value="Clear Chat History")

        with gr.Column():
            new_chat_button = gr.ClearButton(value="New Chat")

    # Sample questions
    with gr.Row():
        with gr.Column():
            gr.Examples([
                ["Explain these documents to me in simpler terms."],
                ["What does these documents talk about?"],
                ["Give the key topics covered in these documents in less than 10 words."],
                ["What are the key findings in these documents?"],
            ], inputs=msg, label= "Click on any example to copy in the chatbox"
            )

  
    # Get a response when a message is submitted to the chatbot
    msg.submit(
        fn = respond, 
        inputs = [msg, chatbot], 
        outputs = [msg, chatbot],
        queue = True)
    

    # Clear the chat history/ New chat
    clear_history_button.click(lambda: [None, None], outputs=[msg, chatbot])
    new_chat_button.click(
        lambda: [None, None, None, None, DEFAULT_STATUS], 
        outputs=[msg, chatbot, pdf_files_dnd, status_message])

    # Description at the bottom of the application
    gr.HTML(description)

# Enable queing
chat.queue()
chat.launch()