import time
import uuid
import pandas as pd
import gradio as gr
from llama_index import GPTSimpleVectorIndex, MockLLMPredictor, ServiceContext

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

llm_predictor = MockLLMPredictor()
service_context_mock = ServiceContext.from_defaults(llm_predictor=llm_predictor)

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

def find_doc(opt, msg2):
    message = ""
    if len(msg2.strip()) < 1:
        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:
            resp = index.query(msg2, service_context=service_context_mock)
            for key, item in resp.extra_info.items():
                message += f"Document: {key}\nExtra details:\n"
                for sub_key, sub_item in item.items():
                    message += f"---- {sub_key}: {sub_item}"

        except Exception as e:
            message = "An error occured when handling your query, please try again."
            print(e)
    return message, ""

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
        """)
        for ex in examples:
            btn = gr.Button(ex)
            btn.click(respond2, [btn, chatbot, msg], [btn, chatbot, msg])

    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])
    gr.Markdown(
    """

    ## Search the matching document
    """)
    opt = gr.Textbox(
        show_label=False,
        placeholder="The document matching with your query will be shown here.",
        interactive=False,
        lines=8
    )
    with gr.Row():
        with gr.Column(scale=0.85):
            msg2 = 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_send2 = gr.Button("Send your query")

    btn_send2.click(find_doc, [opt, msg2], [opt, msg2])
    
if __name__ == "__main__":
    demo.launch()