import gradio as gr
import zipfile
import json
import os
import tempfile
import random
from PIL import Image
import io
import base64

def load_dataset(zip_file):
    global temp_dir
    temp_dir = tempfile.mkdtemp()
    
    try:
        with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        json_file = next(f for f in os.listdir(temp_dir) if f.endswith('.json'))
        with open(os.path.join(temp_dir, json_file), 'r') as f:
            dataset = json.load(f)
        
        return dataset, "Dataset loaded successfully."
    except Exception as e:
        return None, f"Error loading dataset: {str(e)}"

def view_images(zip_file, shuffle, page=1, images_per_page=50):
    global dataset
    
    if zip_file is None:
        return None, "Please upload a dataset ZIP file."
    
    if dataset is None:
        dataset, message = load_dataset(zip_file)
        if dataset is None:  # Error occurred
            return None, message
    
    if shuffle:
        random.shuffle(dataset)
    
    start_idx = (page - 1) * images_per_page
    end_idx = start_idx + images_per_page
    page_data = dataset[start_idx:end_idx]
    
    images = []
    for item in page_data:
        img_path = os.path.join(temp_dir, "images", item['image'])
        img = Image.open(img_path)
        temp_img_path = os.path.join(temp_dir, f"temp_{item['image']}")
        img.save(temp_img_path, format="PNG")
        images.append((temp_img_path, item['tags']))
    
    total_pages = (len(dataset) + images_per_page - 1) // images_per_page
    status = f"Showing page {page} of {total_pages}"
    
    return images, status

def next_page(zip_file, shuffle, current_page):
    return view_images(zip_file, shuffle, page=current_page + 1)

def prev_page(zip_file, shuffle, current_page):
    return view_images(zip_file, shuffle, page=max(1, current_page - 1))

temp_dir = None
dataset = None

with gr.Blocks(theme=gr.themes.Soft()) as iface:
    gr.Markdown("# Dataset Viewer")
    gr.Markdown("Upload a dataset ZIP file to view its images and tags.")
    
    with gr.Row():
        zip_file = gr.File(label="Upload Dataset ZIP", file_types=[".zip"])
        shuffle_checkbox = gr.Checkbox(label="Shuffle Images", value=False)
    
    view_button = gr.Button("View Images")
    
    gallery = gr.Gallery(label="Images", show_label=False, elem_id="gallery", columns=3, rows=3, height="auto")
    status_text = gr.Textbox(label="Status", interactive=False)
    
    with gr.Row():
        prev_button = gr.Button("Previous Page")
        next_button = gr.Button("Next Page")
    
    current_page = gr.State(value=1)
    
    view_button.click(
        view_images,
        inputs=[zip_file, shuffle_checkbox],
        outputs=[gallery, status_text],
        show_progress=True,
    ).then(
        lambda: 1,
        outputs=current_page
    )
    
    next_button.click(
        next_page,
        inputs=[zip_file, shuffle_checkbox, current_page],
        outputs=[gallery, status_text],
        show_progress=True,
    ).then(
        lambda x: x + 1,
        inputs=current_page,
        outputs=current_page
    )
    
    prev_button.click(
        prev_page,
        inputs=[zip_file, shuffle_checkbox, current_page],
        outputs=[gallery, status_text],
        show_progress=True,
    ).then(
        lambda x: max(1, x - 1),
        inputs=current_page,
        outputs=current_page
    )
    
    shuffle_checkbox.change(
        view_images,
        inputs=[zip_file, shuffle_checkbox],
        outputs=[gallery, status_text],
        show_progress=True,
    ).then(
        lambda: 1,
        outputs=current_page
    )

iface.launch()