import json
import requests
import gradio as gr

STYLE = """
#chatbot {
    height: 700px !important;
}
"""

def download_fn(url):
    resp = requests.get(url)
    data = json.loads(resp.content)
    print(len(data))

    for idx, data_item in enumerate(data[:]):
        for conv in data_item:
            if len(conv) != 2:
                del data[idx]
                break

    first_data = []
    if len(data) > 0:
        for first_conv in data[0]:
            piece = []
            for key in first_conv:
                piece.append(first_conv[key])

            first_data.append(piece)
        
    return (
        data,
        first_data,
        0,
        gr.Button(interactive=True) if len(data) > 0 else gr.Button(interactive=False),
        gr.Button(interactive=True) if len(data) > 0 else gr.Button(interactive=False)
    )

def move_next_fn(data, cursor):
    cursor = cursor + 1
    next_data = []

    for conv in data[cursor]:
        piece = []
        for key in conv:
            piece.append(conv[key])

        next_data.append(piece)

    return (
        next_data,
        cursor,
        gr.Button(interactive=True),
        gr.Button(interactive=True) if cursor < len(data)-1 else gr.Button(interactive=False)
    )

def move_prev_fn(data, cursor):
    cursor = cursor -1
    prev_data = []

    for conv in data[cursor]:
        piece = []
        for key in conv:
            piece.append(conv[key])

        prev_data.append(piece)

    return (
        prev_data,
        cursor,
        gr.Button(interactive=True) if cursor > 0 else gr.Button(interactive=False),
        gr.Button(interactive=True)
    )

with gr.Blocks(css=STYLE) as demo:
    data = gr.State([])
    data_to_be_removed = gr.State(set())
    cursor = gr.State(-1)

    with gr.Row():
        url = gr.Textbox(label="URL", placeholder="URL to download conversation recorded JSON", scale=9)
        download = gr.Button("Download", scale=1)
    
    chatbot = gr.Chatbot(elem_id="chatbot")

    with gr.Row():
        prev_btn = gr.Button("Prev", interactive=False)
        mark_btn = gr.Button("✅ Mark to remove ", interactive=False)
        next_btn = gr.Button("Next", interactive=False)

    with gr.Row():
        cleanup_btn = gr.Button("Clean-up")
        get_btn = gr.Button("Download clean-up version")

    download.click(
        download_fn,
        inputs=[url], 
        outputs=[data, chatbot, cursor, next_btn, mark_btn]
    )

    prev_btn.click(
        move_prev_fn,
        inputs=[data, cursor],
        outputs=[chatbot, cursor, prev_btn, next_btn]
    )

    next_btn.click(
        move_next_fn,
        inputs=[data, cursor],
        outputs=[chatbot, cursor, prev_btn, next_btn]
    )

    mark_btn.click(
        lambda remove_list, cursor: remove_list.add(cursor),
        inputs=[data_to_be_removed, cursor],
        outputs=[data_to_be_removed]
    )

    

demo.launch()