import gradio as gr
from PIL import Image, PngImagePlugin
import io

def extract_metadata(image):
    """Extracts Stable Diffusion metadata from a PIL Image."""
    if image:
        # Access metadata using the 'info' dictionary
        metadata = image.info.get("parameters", "")
    else:
        metadata = ""
    return image, metadata

def update_metadata(image, new_metadata):
    """Updates the Stable Diffusion metadata in a PIL Image."""
    if image:
        # Create a PngInfo object to store metadata
        pnginfo = PngImagePlugin.PngInfo()
        pnginfo.add_text("parameters", new_metadata)

        # Save the image to a BytesIO object with the updated metadata
        output_bytes = io.BytesIO()
        image.save(output_bytes, format="PNG", pnginfo=pnginfo)
        output_bytes.seek(0)

        # Re-open the image from the BytesIO object to ensure Gradio displays the updated image
        updated_image = Image.open(output_bytes)

        return updated_image, new_metadata # Return updated PIL Image
    else:
        return None, ""

def clear_metadata(image, metadata):
    """Clears the Stable Diffusion Metadata."""
    return update_metadata(image, "")


def copy_metadata(metadata: str):
    """Copies the metadata to the clipboard (simulated in Gradio)."""
    metadata = metadata.strip()
    embed_index = metadata.find("\nUsed embeddings: ")
    if embed_index < 0:
        embed_index = metadata.find("\nUsed embeddings: ")  # Corrected typo
    if embed_index > 0:
        metadata = metadata[:embed_index]

    return metadata

def process_image(image):
    """Combines extraction and initial display."""
    return extract_metadata(image)


with gr.Blocks() as demo:
    gr.Markdown("# SD EXIF Editor (Gradio)")
    gr.Markdown("Upload a PNG image generated by Stable Diffusion WebUI (AUTOMATIC1111) to view and edit its metadata.")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload PNG Image")
            image_output = gr.Image(type="pil", label="Processed Image") # Keep as PIL

        with gr.Column():
            metadata_textbox = gr.Textbox(label="Metadata", lines=5)
            with gr.Row():
                copy_button = gr.Button("Copy Metadata")
                clear_button = gr.Button("Clear Metadata")
                save_button = gr.Button("Save Metadata")

    # --- Event Handling ---
    image_input.change(process_image, inputs=[image_input], outputs=[image_output, metadata_textbox])
    save_button.click(update_metadata, inputs=[image_input, metadata_textbox], outputs=[image_output, metadata_textbox])  # Use image_input
    clear_button.click(clear_metadata, inputs=[image_input, metadata_textbox], outputs=[image_output, metadata_textbox])  # Use image_input
    copy_button.click(copy_metadata, inputs=[metadata_textbox], outputs=[metadata_textbox])

demo.launch(share=True) # Enable public link