import json
import skia
import gradio as gr
import os
import random

# Define the category dictionary
category_dict = {
    0: "Algorithm",
    1: "Caption",
    2: "Equation",
    3: "Figure",
    4: "Footnote",
    5: "List",
    7: "Table",
    8: "Text",
    9: "Text-EQ",
    10: "Title",
    12: "PaperTitle",
    13: "Code",
    14: "Abstract"
}

# Function to draw bounding box and label on the image using Skia
def draw_annotation(canvas, bbox, category, source_code):
    paint = skia.Paint(Color=skia.ColorRED, Style=skia.Paint.kStroke_Style, StrokeWidth=2)
    text_paint = skia.Paint(Color=skia.ColorRED)
    # text_paint.setTextSize(12)
    
    # Unpack the bounding box coordinates
    x_min, y_min, x_max, y_max = bbox
    
    # Draw the bounding box
    canvas.drawRect(skia.Rect.MakeLTRB(x_min, y_min, x_max, y_max), paint)
    
    # Draw the category label
    label = f"{category_dict[category]}: {source_code}"
    canvas.drawSimpleText(label, x_min, y_min - 10, skia.Font(None, 12), text_paint)

# Load the annotations from a JSON file
def load_annotations(file_path):
    with open(file_path, 'r') as file:
        annotations = json.load(file)
    return annotations

# Main function to render annotations onto images using Skia
def render_annotations(annotations, image_dir):
    annotated_images = []
    for annotation in annotations:
        page_index = annotation['page_index']
        image_path = f"{image_dir}/page_{page_index:04d}.jpg"
        
        # Load the image using Skia
        image = skia.Image.MakeFromEncoded(skia.Data.MakeFromFileName(image_path))
        canvas = skia.Surface(image.width(), image.height()).getCanvas()
        canvas.drawImage(image, 0, 0)
        
        # Draw the annotation on the image
        draw_annotation(canvas, annotation['bbox'], annotation['category'], annotation['source_code'])
        
        # Save the annotated image to a buffer
        output_path = f"annotated_page_{page_index:04d}.png"
        surface = canvas.getSurface()
        # image_buffer = surface.makeImageSnapshot().encode(skia.kPNG_Codec)
        # annotated_images.append((image_buffer, f"Page {page_index}"))
        surface.makeImageSnapshot().save(output_path)
        annotated_images.append(output_path)
    
    return annotated_images

# Gradio interface function
def gradio_interface():
    example_dir = "examples"
    random_index = random.choice([i+1 for i in range(len(os.listdir(example_dir)))])
    selected_example_dir_path = os.path.join(example_dir, str(random_index))
    # annotations_file = random.choice([f for f in os.listdir(example_dir) if f.endswith('.json')])
    annotations_file_path = os.path.join(selected_example_dir_path, "reading_annotation.json")
    annotations = load_annotations(annotations_file_path)
    annotated_images = render_annotations(annotations, selected_example_dir_path)
    return annotated_images

# Create Gradio interface
iface = gr.Interface(
    fn=gradio_interface,
    inputs=[],
    outputs=gr.Gallery(label="Annotated Images"),
    title="Paper Annotation Renderer",
    description="Click to randomly choose an example and render annotations onto the images."
)

# Launch the Gradio interface
iface.launch()