import gradio as gr
import json
import os
import zipfile
from datetime import datetime

# Initialize dataset storage
datasets = {}

# Function to save dataset to JSONL file
def save_dataset(dataset_name):
    if dataset_name in datasets:
        dataset = datasets[dataset_name]
        os.makedirs(f"{dataset_name}/images", exist_ok=True)
        with open(f"{dataset_name}/dataset.jsonl", "w") as f:
            for item in dataset:
                f.write(json.dumps(item) + "\n")

# Function to load dataset from JSONL file
def load_dataset(dataset_name, file):
    if file:
        os.makedirs(dataset_name, exist_ok=True)
        with open(f"{dataset_name}/dataset.jsonl", "r") as f:
            datasets[dataset_name] = [json.loads(line) for line in f]
        return f"Dataset '{dataset_name}' loaded successfully."
    return "No file uploaded."

# Function to add image/prompt pair to dataset
def add_to_dataset(dataset_name, image, prompt):
    if dataset_name not in datasets:
        datasets[dataset_name] = []
    image_path = f"{dataset_name}/images/{datetime.now().timestamp()}.png"
    image.save(image_path)
    datasets[dataset_name].append({"image": image_path, "prompt": prompt})
    return f"Added to dataset '{dataset_name}'."

# Function to edit image/prompt pair in dataset
def edit_dataset(dataset_name, index, new_prompt):
    if dataset_name in datasets and 0 <= index < len(datasets[dataset_name]):
        datasets[dataset_name][index]["prompt"] = new_prompt
        return f"Edited prompt at index {index} in dataset '{dataset_name}'."
    return "Invalid index."

# Function to delete image/prompt pair from dataset
def delete_from_dataset(dataset_name, index):
    if dataset_name in datasets and 0 <= index < len(datasets[dataset_name]):
        os.remove(datasets[dataset_name][index]["image"])
        datasets[dataset_name].pop(index)
        return f"Deleted item at index {index} from dataset '{dataset_name}'."
    return "Invalid index."

# Function to generate and download dataset as zip file
def generate_dataset(dataset_name):
    if dataset_name in datasets:
        save_dataset(dataset_name)
        zip_filename = f"{dataset_name}.zip"
        with zipfile.ZipFile(zip_filename, 'w') as zipf:
            for root, _, files in os.walk(dataset_name):
                for file in files:
                    zipf.write(os.path.join(root, file))
        return zip_filename
    return "Dataset not found."

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Dataset Builder")
    
    with gr.Tab("Create/Edit Dataset"):
        dataset_name = gr.Textbox(label="Dataset Name")
        image = gr.Image(label="Upload Image")
        prompt = gr.Textbox(label="Prompt")
        add_button = gr.Button("Add to Dataset")
        add_output = gr.Textbox()
        add_button.click(add_to_dataset, inputs=[dataset_name, image, prompt], outputs=add_output)
        
        index = gr.Number(label="Index to Edit/Delete", precision=0)
        new_prompt = gr.Textbox(label="New Prompt")
        edit_button = gr.Button("Edit Prompt")
        edit_output = gr.Textbox()
        edit_button.click(edit_dataset, inputs=[dataset_name, index, new_prompt], outputs=edit_output)
        
        delete_button = gr.Button("Delete from Dataset")
        delete_output = gr.Textbox()
        delete_button.click(delete_from_dataset, inputs=[dataset_name, index], outputs=delete_output)
    
    with gr.Tab("Upload Dataset"):
        upload_dataset_name = gr.Textbox(label="Dataset Name")
        upload_file = gr.File(label="Upload JSONL File")
        upload_button = gr.Button("Upload Dataset")
        upload_output = gr.Textbox()
        upload_button.click(load_dataset, inputs=[upload_dataset_name, upload_file], outputs=upload_output)
    
    with gr.Tab("Generate/Download Dataset"):
        generate_button = gr.Button("Generate Dataset")
        generate_output = gr.File()
        generate_button.click(generate_dataset, inputs=[dataset_name], outputs=generate_output)
    
    with gr.Tab("View Dataset"):
        view_dataset_name = gr.Textbox(label="Dataset Name")
        view_button = gr.Button("View Dataset")
        view_output = gr.JSON()
        view_button.click(lambda name: datasets.get(name, []), inputs=[view_dataset_name], outputs=view_output)

demo.launch()