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 = {}
        if 'metadata' in image.info:
            metadata = json.loads(image.info['metadata'])
        elif 'prompt' in image.info:
            metadata = json.loads(image.info['prompt'])
        elif 'Comment' in image.info:
            metadata = json.loads(image.info['Comment'])
            metadata['model'] = 'NovelAI'
        elif 'parameters' in image.info:
            if image.info['parameters'].startswith('{'):
                parameters_data = json.loads(image.info['parameters'])
                if 'sui_image_params' in parameters_data:
                    sui_image_params = parameters_data['sui_image_params']
                    metadata.update(sui_image_params)
                else:
                    metadata = parameters_data
            else:
                lines = image.info['parameters'].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 line.startswith('Steps:'):
                        steps_info = line.split(':', 1)[1].strip().split(',')
                        for info in steps_info:
                            info = info.strip()
                            if ':' in info:
                                key, value = info.split(':', 1)
                                metadata[key.strip()] = value.strip()
        else:
            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")

    input_image.change(
        fn=process_image,
        inputs=input_image,
        outputs=[status_output, output_metadata],
        api_name="interrogate"
    )

    gr.Examples(
        examples=[
            ["example/asukatest.png"],
            ["example/arimakana.png"],
            ["example/stelle.png"],
            ["example/shinji.png"],
        ],
        inputs=input_image,
        outputs=[status_output, output_metadata],
        fn=process_image,
        cache_examples=True,
    )

demo.launch()