import gradio as gr
import random
import time
from agent_t5 import Agent
from config import Config
from retrieval.retrieval import BM25


args = Config()
chatbot = Agent(args)
answer_areas = list(range(args.choices))
context_areas = list(range(args.choices))

with gr.Blocks() as demo:
    # gr.Markdown("Flip text or image files using this demo.")
    with gr.Tab("Chatbot"):
        with gr.Row():
            with gr.Column():
                chatbot_area = gr.Chatbot().style(height=700)
                msg = gr.Textbox(label='Your prompt')

            with gr.Column(scale=0.15, min_width=300):
                for i in range(args.choices):
                    with gr.Accordion(f"Answer: {i+1}", open=False) as answer_areas[i]:
                        context_areas[i] = gr.Markdown(f"Context {i+1}") 
                clear_chat = gr.Button("Clear history")

    with gr.Tab("Your context"):
        context_box = gr.Textbox(
            label='Your context here! You can upload a context file or typing context here and click "Using context"',
            lines=20,
            placeholder="Enter your context here..."
        )
        
        with gr.Row() as taskbar:
            upload_btt = gr.UploadButton('Upload Context File')
            clear_context_btt = gr.Button("Clear context")
            context_btt = gr.Button("Using context")


    def user(user_message, history):
        print("Context box value:", context_box.info)
        return "", history + [[user_message, None]]

    def bot(history):
        question = history[-1][0]
        print('User mess:', question)
        answers, contexts = chatbot.asking(question)

        for i in range(chatbot.choices):
            context_areas[i].value = contexts[i]['context']
            answer_areas[i].value = answers[i]


        print(answers)
        history[-1][1] = ""
        for character in answers[0]:
            history[-1][1] += character
            time.sleep(0.01)
            yield history

    msg.submit(user, [msg, chatbot_area], [msg, chatbot_area], queue=False).then(
        bot, chatbot_area, chatbot_area
    )

    context_btt.click(chatbot.get_context, [context_box, ])
    upload_btt.upload(chatbot.load_context, [upload_btt, ], context_box)
    clear_context_btt.click(chatbot.clear_context, outputs=context_box)

demo.queue()
demo.launch()