import gradio as gr
from argparse import ArgumentParser
import time
from pathlib import Path


def add_message(history, mic, text):
    if not mic and not text:
        return history, "Input is empty"
    if mic and Path(mic).exists():
        history.append({"role": "user", "content": {"path": mic}})
    if text:
        history.append({"role": "user", "content": text})
    print(f"{history=}")
    return history, None


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


def regenerate(task_history):
    """Regenerate the last bot response."""
    while task_history and task_history[-1]["role"] == "assistant":
        print(f"discard {task_history[-1]}")
        task_history.pop()
    if task_history:
        print(f"{task_history=}")
        return predict(task_history)


def predict(chatbot):
    """Generate a response from the model."""
    if len(chatbot) > 0:
        for item in reversed(chatbot):
            if item["role"] != "user":
                break
            if isinstance(item["content"], dict):
                chatbot.append({"role": "assistant", "content": item["content"]})
            elif isinstance(item["content"], tuple):
                chatbot.append({"role": "assistant", "content": {"path":item["content"][0]}})
            else:
                chatbot.append({"role": "assistant", "content": f"echo {item['content']}"})
    else:
        response = "**That's cool!**"
        chatbot.append({"role": "assistant", "content": response})
    return 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"
        )
        # task_history = gr.State([])
        mic = gr.Audio(type="filepath")
        text = gr.Textbox(placeholder="Enter message ...")

        with gr.Row():
            clean_btn = gr.Button("🧹 Clean History (清除历史)")
            regen_btn = gr.Button("🤔️ Regenerate (重试)")
            submit_btn = gr.Button("🚀 Submit")

        def on_submit(chatbot, mic, text):
            history, error = add_message(chatbot, mic, text)
            if error:
                gr.Warning(error)  # 显示警告消息
                return history, None, None
            else:
                return predict(history), None, None

        submit_btn.click(
            fn=on_submit,
            inputs=[chatbot, mic, text],
            outputs=[chatbot, mic, text],
        )
        clean_btn.click(reset_state, outputs=[chatbot], show_progress=True)
        regen_btn.click(regenerate, [chatbot], [chatbot], show_progress=True)

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


if __name__ == "__main__":
    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="0.0.0.0", help="Demo server name."
    )

    args = parser.parse_args()

    _launch_demo(args)