import time
import pandas as pd
import gradio as gr
from llama_index import GPTSimpleVectorIndex

title = "Confidential forensics tool with ChatGPT"
examples = ["Who is Phillip Allen?", "What is the project in Austin is about?", "Give me more details about the real estate project"]

index = GPTSimpleVectorIndex.load_from_disk('email.json')
docs_arr = []
for doc in index.docstore.docs:
    docs_arr.append(doc)
dat_fr = pd.DataFrame({"Documents loaded": docs_arr})

def respond_upload(btn_upload, message, chat_history):
    time.sleep(2)
    message = "***File uploaded***"
    bot_message = "Your document has been uploaded and will be accounted for your queries."
    chat_history.append((message, bot_message))
    return btn_upload, "", chat_history

def respond2(message, chat_history, box, btn):
    if len(message.strip()) < 1:
        message = "***Empty***"
        bot_message = "Oops, it looks like your query was not valid. Please make sure you typed something in your text box and then try again."
    else:
        try:
            bot_message = str(index.query(message)).strip()
        except:
            bot_message = "An error occured when handling your query, please try again."
    chat_history.append((message, bot_message))
    return message, chat_history, box

def respond(message, chat_history):
    if len(message.strip()) < 1:
        message = "***Empty***"
        bot_message = "Oops, it looks like your query was not valid. Please make sure you typed something in your text box and then try again."
    else:
        try:
            bot_message = str(index.query(message)).strip()
        except:
            bot_message = "An error occured when handling your query, please try again."
    chat_history.append((message, bot_message))
    return "", chat_history

with gr.Blocks(title=title) as demo:
    gr.Markdown(
    """
    # """ + title + """
    ...
    """)
    dat = gr.Dataframe(
            value=dat_fr
    )
    gr.Markdown(
    """
    ## Chatbot
    """)
    chatbot = gr.Chatbot().style(height=400)
    with gr.Row():
        with gr.Column(scale=0.70):
            msg = gr.Textbox(
                show_label=False,
                placeholder="Enter text and press enter, or click on Send.",
            ).style(container=False)
        with gr.Column(scale=0.15, min_width=0):
            btn_send = gr.Button("Send your query")
        with gr.Column(scale=0.15, min_width=0):
            btn_upload = gr.UploadButton("Upload a new document...", file_types=["text"])
    with gr.Row():
        gr.Markdown(
        """
        Example of queries
        """)
        i = 0
        for ex in examples:
            btn = gr.Button(examples[i])
            btn.click(respond2, [btn, chatbot, msg], [btn, chatbot, msg])
            i += 1

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    btn_send.click(respond, [msg, chatbot], [msg, chatbot])
    btn_upload.upload(respond_upload, [btn_upload, msg, chatbot], [btn_upload, msg, chatbot])

if __name__ == "__main__":
    demo.launch()