import gradio as gr
from argparse import ArgumentParser
import time


def _get_args():
    parser = ArgumentParser()
    parser.add_argument("--inbrowser", action="store_true", default=False,
                        help="Automatically launch the interface in a new tab on the default browser.")
    parser.add_argument("--server-port", type=int, default=7860,
                        help="Demo server port.")
    parser.add_argument("--server-name", type=str, default="127.0.0.1",
                        help="Demo server name.")

    args = parser.parse_args()
    return args


def add_message(history, message):
    for x in message["files"]:
        history.append({"role": "user", "content": {"path": x}})
    if message["text"] is not None:
        history.append({"role": "user", "content": message["text"]})
    return history, gr.MultimodalTextbox(value=None, interactive=False)

def reset_user_input():
    """Reset the user input field."""
    return gr.Textbox.update(value='')


def reset_state(task_history):
    """Reset the chat history."""
    return [], []


def regenerate(task_history):
    """Regenerate the last bot response."""
    if task_history and task_history[-1]['role'] == 'assistant':
        task_history.pop()
    if task_history:
        task_history = yield from predict(task_history)

def predict(chatbot):
    """Generate a response from the model."""
    if len(chatbot) > 0:
        response = f"chatbot test: {chatbot[-1]}"
    else:
        response = "**That's cool!**"
    chatbot.append({"role": "assistant", "content": ""})
    for character in response:
        chatbot[-1]["content"] += character
        time.sleep(0.05)
        yield chatbot


def _launch_demo(args):
    with gr.Blocks() as demo:
        gr.Markdown("""<center><font size=8>ChatBot Dummy</center>""")
        chatbot = gr.Chatbot(elem_id="chatbot", bubble_full_width=False, type="messages")

        user_input = gr.MultimodalTextbox(
            interactive=True,
            placeholder="Enter message or upload file...",
            show_label=False,
            file_types=["audio"],
            # sources=["microphone", "upload"],
        )
        # task_history = gr.State([])

        with gr.Row():
            empty_bin = gr.Button("🧹 Clear History (清除历史)")
            regen_btn = gr.Button("🤔️ Regenerate (重试)")

        user_input.submit(fn=add_message,
                          inputs=[chatbot, user_input],
                          outputs=[chatbot, user_input],concurrency_limit = 4).then(
            predict, [chatbot], [chatbot], show_progress=True
        ).then(lambda: gr.MultimodalTextbox(interactive=True), None, [user_input])
        empty_bin.click(reset_state, outputs=[chatbot], show_progress=True, concurrency_limit = 4)
        regen_btn.click(regenerate, [chatbot], [chatbot], show_progress=True, concurrency_limit = 4)

    demo.queue().launch(
        share=False,
        inbrowser=args.inbrowser,
        server_port=args.server_port,
        server_name=args.server_name,
        max_threads=4
    )


if __name__ == "__main__":
    args = _get_args()
    _launch_demo(args)