import gradio as gr
from PIL import Image, PngImagePlugin
import json
import traceback

def extract_metadata(image):
    if image is None:
        return "Please upload an image.", {}

    try:
        metadata = {}
        # Handling multiple possible metadata keys
        potential_keys = ['metadata', 'prompt', 'Comment', 'parameters', 'exif']
        for key in potential_keys:
            if key in image.info:
                raw_data = image.info[key]
                
                # If raw data starts with '{', assume JSON format
                if raw_data.startswith('{'):
                    metadata = json.loads(raw_data)
                else:
                    if key == 'parameters':
                        # Attempt to process Stable Diffusion or NovelAI style data
                        lines = raw_data.split('\n')
                        prompt = lines[0].strip()
                        negative_prompt = lines[1].strip().replace('Negative prompt:', '').strip()
                        metadata['prompt'] = prompt
                        metadata['negative_prompt'] = negative_prompt

                        for line in lines[2:]:
                            line = line.strip()
                            if ':' in line:
                                key, value = line.split(':', 1)
                                metadata[key.strip()] = value.strip()
                    elif key == 'Comment':
                        # Specific handling for NovelAI
                        metadata = json.loads(raw_data)
                        metadata['model'] = 'NovelAI'
                break  # Exit loop once a supported key is found

        if not metadata:
            return "No supported metadata found in the image.", {}

        return "Metadata extracted successfully.", metadata
    except Exception as e:
        error_message = f"Error extracting metadata: {str(e)}\n{traceback.format_exc()}"
        return error_message, {}

def process_image(image):
    status, metadata = extract_metadata(image)
    return status, metadata

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Image Metadata Extractor
        Extract and display metadata from images generated by various AI tools.
        """
    )

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image", type="pil", height=480)

        with gr.Column():
            status_output = gr.Textbox(label="Status")
            output_metadata = gr.JSON(label="Metadata")

    # Event listener for when the image is changed
    input_image.change(
        fn=process_image,
        inputs=input_image,
        outputs=[status_output, output_metadata],
        api_name="interrogate"
    )

demo.launch()