import gradio as gr

from utils.chatbot import Chatbot, VideoChatbot
from utils.utils import *

def create_new_model():
    return Chatbot()

# start of gradio interface
with gr.Blocks() as demo:
    user_chatbot = gr.State(create_new_model)
    test_video_chabot = gr.State(VideoChatbot())

    with gr.Row():
        gr.HTML("Junyi Academy Chatbot")
    with gr.Row(equal_height=True):
        with gr.Column(scale=5):
            with gr.Row():
                chatbot = gr.Chatbot()

            with gr.Row():
                with gr.Column(scale=12):
                    user_input = gr.Textbox(
                        show_label=False,
                        placeholder="Enter text",
                        container=False,
                    )

                with gr.Column(min_width=70, scale=1):
                    clear_btn = gr.Button("清除")
                with gr.Column(min_width=70, scale=1):
                    submit_btn = gr.Button("傳送")

    with gr.Row():
        index_file = gr.File(
            file_count="multiple", file_types=["pdf"], label="Upload PDF file", scale=3
        )
        upload_to_db = gr.CheckboxGroup(
            ["Upload to Database"],
            label="是否上傳至資料庫", info="須在上傳檔案前勾選（將資料上傳至資料庫時，資料庫會自動建立索引，下次使用時可以直接檢索，無勾選則僅作這次使用）", scale=1
        )

    with gr.Row():
        instruction = gr.Markdown(
            """
            ## 使用說明
            1. 上傳一個或多個 PDF 檔案，系統將自動進行摘要、翻譯等處理後建立知識庫
            2. 在上方輸入欄輸入問題，系統將自動回覆
            3. 可以根據下方的摘要內容來提問
            4. 每次對話會根據第一個問題的內容來檢索所有文件，並挑選最能回答問題的文件來回覆
            5. 要切換檢索的文件，請點選「清除」按鈕後再重新提問

        """,
        )

    with gr.Row():
        describe = gr.Markdown("", visible=True)

    with gr.Row():
        video_text_input = gr.Textbox("", visible=False)
        video_text_output = gr.Textbox("", visible=False)

        transcript_id = gr.Textbox("", visible=False)
        user_question = gr.Textbox("", visible=False)
        content_output = gr.Textbox("", visible=False)

    # end of gradio interface

    # start of workflow controller

    # defining workflow of user bot interaction
    bot_args = dict(
        fn=bot,
        inputs=user_chatbot,
        outputs=chatbot,
    )

    user_args = dict(
        fn=user,
        inputs=[user_chatbot, user_input],
        outputs=[user_input, chatbot],
        queue=False,
    )

    response = user_input.submit(**user_args).then(**bot_args)

    response.then(lambda: gr.update(interactive=True), None, [user_input], queue=False)

    submit_btn.click(
        **user_args,
    ).then(
        **bot_args
    ).then(lambda: gr.update(interactive=True), None, [user_input], queue=False)


    # defining workflow of clear state
    clear_state_args = dict(
        fn=clear_state,
        inputs=user_chatbot,
        outputs=None,
    )

    clear_btn.click(**clear_state_args)

    # defining workflow of building knowledge base
    send_system_nofification_args = dict(
        fn=send_system_nofification,
        inputs=user_chatbot,
        outputs=chatbot,
    )

    bulid_knowledge_base_args = dict(
        fn=build_knowledge_base,
        inputs=[user_chatbot, index_file, upload_to_db],
        outputs=None,
    )

    change_md_args = dict(
        fn=change_md,
        inputs=[user_chatbot],
        outputs=[describe],
    )

    index_file.upload(**send_system_nofification_args).then(
        lambda: gr.update(interactive=True), None, None, queue=False
    ).then(**bulid_knowledge_base_args).then(**send_system_nofification_args).then(
        lambda: gr.update(interactive=True), None, None, queue=False
    ).then(
        **change_md_args
    )

    video_text_input.submit(video_bot, [test_video_chabot, video_text_input], video_text_output, api_name="video_bot")
    transcript_id.submit(search_transcript_content, [transcript_id, user_question], content_output, api_name="search_transcript_content")


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