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)
RLHF = []

with gr.Blocks() as demo:
    gr.Markdown( """<h1><center>Vietnamese QA assistant</center></h1>""")
    with gr.Tab("Chatbot"):
        with gr.Row():
            with gr.Column():
                chatbot_area = gr.Chatbot(
                    [[None, "Xin chào!"],
                     [None, "Tôi là một công cụ đọc hiểu cho tiếng Việt. Bạn có thể đặt bất kì câu hỏi cho tôi hoặc bạn có thể upload văn bản của bạn và hỏi tôi về nó ở tab 'Your context'"],
                     ["Công dụng của paracetamol", "hạ sốt, giảm đau"]]
                ).style(height=600)
                msg = gr.Textbox(label='Your prompt')

            with gr.Column(scale=0.3, min_width=500):
                for i in range(args.show_c):
                    with gr.Accordion(f"Answer {i+1}", open=False) as answer_area:
                        with gr.Row():
                            gr.Button('👍 Answer').style(size='sm', full_width=False)
                            gr.Button('👎🏻 Answer').style(size='sm', full_width=False)
                        context_area = gr.Markdown(f"Context {i+1}")
                        # gr.CheckboxGroup(choices=['Good Answer', 'Bad Answer', 'Good Context', 'Bad Context'])

                        with gr.Row():
                            gr.Button('👍 Context').style(size='sm', full_width=False)
                            gr.Button('👎🏻 Context').style(size='sm', full_width=False)

                        RLHF.append(answer_area)
                        RLHF.append(context_area)
                clear_chat = gr.Button("Clear history")

    with gr.Tab("Your context"):
        context_box = gr.Textbox(
            label="Bạn có thể upload hoặc viết văn bản của bạn vào đây và quay lại tab 'Chatbot' để hỏi tôi!",
            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]
        answers = chatbot.asking(question)


        print(answers)
        history[-1][1] = ""
        for character in answers:
            history[-1][1] += character
            time.sleep(0.01)
            yield history
    
    def fill_feedback(prompt):
        results = chatbot.temp
        outptus = []
        for i in range(args.show_c):
            if i < chatbot.choices:
                answer = f"{results[i]['score']}% | {results[i]['answer']}"
                outptus.append(RLHF[2*i].update(label=answer, visible=True, open=False))
                outptus.append(results[i]['context'])
            else:
                outptus.append(RLHF[2*i].update(visible=False))
                outptus.append("No context")

        return outptus

    def clear_history(history):
        history = []
        return history

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

    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)
    clear_chat.click(clear_history, chatbot_area, chatbot_area)

demo.queue()
demo.launch(server_name='0.0.0.0')