import gradio as gr
from datasets import load_dataset
from PIL import Image, ImageDraw
import os

hf_token = os.getenv("HF_TOKEN")
dataset = load_dataset("agentsea/wave-ui-25k", split="train", token=hf_token)
current_index = 0

def draw_bounding_box(image, bbox):
    draw = ImageDraw.Draw(image)
    draw.rectangle(bbox, outline="red", width=2)
    return image

def show_image(index):
    global current_index
    current_index = index
    data = dataset[current_index]
    image = data['image']
    bbox = data['bbox']
    
    image_with_bbox = draw_bounding_box(image, bbox)
    
    image_info = f"""
    \n**Image {current_index + 1} of {len(dataset)}**
    \n**Source**: {data['source']}
    \n**Platform**: {data['platform']}
    \n**Name**: {data['name']}
    \n**Description**: {data['description']}
    \n**Type**: {data['type']}
    \n**OCR**: {data['OCR']}
    \n**Language**: {data['language']}
    \n**Purpose**: {data['purpose']}
    \n**Expectation**: {data['expectation']}
    """
    return image_with_bbox, image_info

def next_image():
    global current_index
    current_index = (current_index + 1) % len(dataset)
    return show_image(current_index)

def previous_image():
    global current_index
    current_index = (current_index - 1) % len(dataset)
    return show_image(current_index)

def change_index(new_index):
    new_index = int(new_index)
    if 0 <= new_index < len(dataset):
        return show_image(new_index)
    else:
        return show_image(current_index), f"Invalid index. Please enter a number between 0 and {len(dataset) - 1}."

shortcut_js = """
<script>
function shortcuts(e) {
    var event = document.all ? window.event : e;
    switch (e.target.tagName.toLowerCase()) {
        case "input":
        case "textarea":
        case "select":
        case "button":
        break;
        default:
        if (e.key === "ArrowRight") {
            document.getElementById("next_btn").click();
        } else if (e.key === "ArrowLeft") {
            document.getElementById("prev_btn").click();
        }
    }
}
document.addEventListener('keydown', shortcuts, false);
</script>
"""

with gr.Blocks(head=shortcut_js) as app:
    gr.Markdown("# WaveUI 25k Dataset Explorer")
    
    with gr.Row():
        image_output = gr.Image(type="pil")
        label_output = gr.Markdown()

    with gr.Row():
        prev_button = gr.Button("Previous", elem_id="prev_btn")
        next_button = gr.Button("Next", elem_id="next_btn")
    
    with gr.Row():
        index_input = gr.Number(label="Go to index (0-based)", value=0)
        go_button = gr.Button("Go", elem_id="go_btn")
    
    prev_button.click(previous_image, outputs=[image_output, label_output])
    next_button.click(next_image, outputs=[image_output, label_output])
    go_button.click(change_index, inputs=[index_input], outputs=[image_output, label_output])
    
    app.load(show_image, inputs=[gr.Number(value=0, visible=False)], outputs=[image_output, label_output])

app.launch()