import json

import gradio as gr

max_rows = 10


def get_row(visible=True):
    with gr.Row().style(equal_height=True):
        with gr.Column():
            do_train = gr.Checkbox(interactive=True, label="Do train on this content?", visible=visible)
            role = gr.Textbox(interactive=True, label="Role", visible=visible)
        content = gr.TextArea(interactive=True, label="Content", visible=visible)
    separator = gr.Markdown("---", visible=visible)
    return do_train, role, content, separator


def variable_outputs(k):
    k = int(k)
    return [gr.update(visible=True)] * k + [gr.update(visible=False)] * (max_rows - k)


def respond(*data):
    data = list(data)
    k = int(data[0])
    data = data[1:]
    do_train_items = data[:max_rows][:k]
    role_items = data[max_rows:max_rows * 2][:k]
    content_items = data[max_rows * 2:][:k]
    data = []
    for do_train, role, content in zip(do_train_items, role_items, content_items):
        data.append(
            {
                "do_train": do_train,
                "role": role,
                "content": content
            }
        )
    output = {
        "conversation": data
    }
    return output


with gr.Blocks() as demo:
    export_button = gr.Button("Export")
    # response = gr.JSON(value={}).style(container=True)
    response = gr.Textbox().style(show_copy_button=True)

    s = gr.Slider(1, max_rows, value=1, step=1, label="How many rows to show:")
    do_train_items = []
    role_items = []
    content_items = []
    separator_items = []
    for i in range(max_rows):
        do_train, role, content, separator = get_row(visible=i < s.value)
        do_train_items.append(do_train)
        role_items.append(role)
        content_items.append(content)
        separator_items.append(separator)

    s.change(variable_outputs, s, do_train_items, queue=False)
    s.change(variable_outputs, s, role_items, queue=False)
    s.change(variable_outputs, s, content_items, queue=False)
    s.change(variable_outputs, s, separator_items, queue=False)

    export_button.click(respond, inputs=[s] + do_train_items + role_items + content_items, outputs=response,
                        queue=False)

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