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("Vietnamese Machine Reading Comprehension in healthcare")
    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 câu hỏi cho tôi về lĩnh vực y tế 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=700)
                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()