import gradio as gr


def get_model_response(txt: str, user_chats: list[str], full_chat_history: list[str]):
    response = f'Hello, you said "{txt}"'
    full_chat_history = full_chat_history + [txt, response]
    user_chats.append(txt)
    responses = [
        (full_chat_history[i], full_chat_history[i + 1])
        for i in range(0, len(full_chat_history) - 1, 2)
    ]
    return responses, full_chat_history, user_chats


def clear_chat(
    full_chat_history: list[str],
    chatbox: list[tuple[str, str]],
    user_chats: list[str],
    cleartill_dropdown_choices: list[str],
    drop_down: str,
):
    if drop_down == "":
        index = 0
    else:
        index = user_chats.index(drop_down)
    return (
        # full_chat_history has user_chat and response chats in a single line.
        full_chat_history[: index * 2],
        chatbox[:index],
        user_chats[:index],
        cleartill_dropdown_choices[:index],
        gr.update(choices=cleartill_dropdown_choices[:index], value=""),
    )


def add_to_dropdown(txt, drop_down_options):
    drop_down_options = drop_down_options + [txt]
    return gr.update(choices=drop_down_options), drop_down_options


drop_down = gr.Dropdown(
    label="Clear From (including)",
    choices=[],
    interactive=True,
    multiselect=False,
)
with gr.Blocks() as demo:
    user_chats = gr.State(value=[])
    full_chat_history = gr.State(value=[])
    cleartill_dropdown_choices = gr.State(value=[])

    chatbox = gr.Chatbot()
    with gr.Row():
        with gr.Column(scale=5):
            user_chat_txt = gr.Textbox(
                placeholder="Say Hi!", label="press Enter to submit"
            )
            user_chat_txt.submit(
                get_model_response,
                inputs=[user_chat_txt, user_chats, full_chat_history],
                outputs=[chatbox, full_chat_history, user_chats],
            )
            user_chat_txt.submit(lambda: gr.update(value=""), [], user_chat_txt)
            user_chat_txt.submit(
                add_to_dropdown,
                inputs=[user_chat_txt, cleartill_dropdown_choices],
                outputs=[drop_down, cleartill_dropdown_choices],
            )
        with gr.Column(scale=1, min_width=250):
            drop_down.render()
            clear_btn = gr.Button("Clear", label="Clear chat history")
            clear_btn.click(
                clear_chat,
                inputs=[
                    full_chat_history,
                    chatbox,
                    user_chats,
                    cleartill_dropdown_choices,
                    drop_down,
                ],
                outputs=[
                    full_chat_history,
                    chatbox,
                    user_chats,
                    cleartill_dropdown_choices,
                    drop_down,
                ],
            )

    with gr.Accordion(label="Red-Teaming", open=True):
        gr.Markdown("*instructions on what to do if you break the model*")
        gr.Textbox()
        with gr.Row():
            submit_btn = gr.Button("Submit", label="Submit")
            import_chat_btn = gr.Button("Import Chat", label="import chat")
            export_chat_btn = gr.Button("Export Chat", label="export chat")

demo.launch()