import gradio as gr
from datasets import load_dataset
from chatgpt_conversation_parser.main import fetch_and_parse_conversation


def add_to_dataset(title, data):
    dataset = load_dataset("beratcmn/parse-gpt-tr", cache_dir="./datasets")
    model = data[0]["model"]["slug"]
    id = data[0]["id"]

    for message in data:
        del message["model"]
        del message["id"]

    data_dict = {"id": id, "model": model, "messages": data}

    dataset = dataset["train"].add_item(data_dict)
    dataset.push_to_hub("beratcmn/parse-gpt-tr")


def action(url, cb):
    result = fetch_and_parse_conversation(url)
    title, data = result
    yield result
    if cb:
        add_to_dataset(title, data)


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # ChatGPT Conversation Parser

        This is a simple tool to fetch and parse a conversation from ChatGPT's web interface.

        Enter the URL of the conversation to fetch and parse. The tool will return a JSONL file and the parsed conversation.
        """.strip()
    )

    gr.Interface(
        fn=action,
        inputs=[
            "text",
            gr.Checkbox(
                True,
                label="Add this conversation to the public dataset?",
            ),
        ],
        outputs=[gr.DownloadButton(label="Download JSONL"), gr.JSON()],
    )

    gr.Markdown(
        """
        Made with ❤️ by [Berat Çimen](https://www.beratcimen.com/).
        """.strip()
    )

    demo.launch()