import gradio as gr
from PIL import Image
from PIL.ExifTags import TAGS
import numpy as np
import re
import json
import tempfile
import os

def get_image_metadata(img):
    """Extract metadata based on image format."""
    metadata = {}
    if img.format == 'PNG':
        metadata = img.info
    elif img.format in ['JPEG', 'TIFF']:
        exif_data = img._getexif()
        if exif_data:
            for tag, value in exif_data.items():
                tag_name = TAGS.get(tag, tag)
                metadata[tag_name] = value
    return metadata

def parse_parameters(parameters):
    """Parse the 'parameters' field to extract prompts and seed number."""
    prompt = 'N/A'
    negative_prompt = 'N/A'
    adetailer_prompt_1 = 'N/A'
    adetailer_prompt_2 = 'N/A'
    seed_number = -1

    # Extract the prompt
    prompt_match = re.search(r'(.*?)Negative prompt:', parameters, re.DOTALL | re.IGNORECASE)
    if prompt_match:
        prompt = prompt_match.group(1).strip()
    else:
        prompt_match = re.search(r'(.*?)(Steps:|$)', parameters, re.DOTALL | re.IGNORECASE)
        if prompt_match:
            prompt = prompt_match.group(1).strip()

    # Extract the negative prompt
    negative_prompt_match = re.search(r'Negative prompt:(.*?)(Steps:|$)', parameters, re.DOTALL | re.IGNORECASE)
    if negative_prompt_match:
        negative_prompt = negative_prompt_match.group(1).strip()

    # Extract ADetailer prompts
    adetailer_prompt_1_match = re.search(r'ADetailer prompt:\s*"(.*?)"', parameters, re.IGNORECASE)
    if adetailer_prompt_1_match:
        adetailer_prompt_1 = adetailer_prompt_1_match.group(1).strip()

    adetailer_prompt_2_match = re.search(r'ADetailer negative prompt 2nd:\s*"(.*?)"', parameters, re.IGNORECASE)
    if adetailer_prompt_2_match:
        adetailer_prompt_2 = adetailer_prompt_2_match.group(1).strip()

    # Extract Seed Number
    seed_match = re.search(r'Seed:\s*(\d+)', parameters, re.IGNORECASE)
    if seed_match:
        seed_number = seed_match.group(1).strip()

    return prompt, negative_prompt, adetailer_prompt_1, adetailer_prompt_2, seed_number

def extract_metadata(image_file):
    """Extract and parse metadata from the uploaded image."""
    try:
        img = Image.open(image_file)
    except Exception as e:
        # Return placeholders in case of error
        return (None, 'Error: Unable to open image.', *(['N/A'] * 7))

    metadata = get_image_metadata(img)
    img_display = np.array(img)

    # Convert metadata to string
    metadata_str = "\n".join([f"{key}: {value}" for key, value in metadata.items()])

    # Initialize default values
    prompt = negative_prompt = adetailer_prompt_1 = adetailer_prompt_2 = 'N/A'
    seed_number = -1
    comfy_prompt = comfy_workflow = 'N/A'

    # Parse parameters if available
    if 'parameters' in metadata:
        parsed = parse_parameters(metadata['parameters'])
        prompt, negative_prompt, adetailer_prompt_1, adetailer_prompt_2, seed_number = parsed

    # Direct extraction with validation for Comfy Prompt
    if 'prompt' in metadata and isinstance(metadata['prompt'], str):
        comfy_prompt = metadata['prompt'].strip()
    else:
        comfy_prompt = 'N/A'  # Fallback if prompt is missing or not a string

    # Direct extraction with validation for Comfy Workflow
    if 'workflow' in metadata and isinstance(metadata['workflow'], str):
        comfy_workflow = metadata['workflow'].strip()
    else:
        comfy_workflow = 'N/A'  # Fallback if workflow is missing or not a string

    return (
        img_display, 
        prompt, 
        negative_prompt, 
        seed_number, 
        adetailer_prompt_1, 
        adetailer_prompt_2, 
        metadata_str, 
        comfy_prompt, 
        comfy_workflow
    )

def export_workflow(workflow_text):
    """
    Converts the workflow text into JSON format and returns it as a downloadable file.
    """
    print("Export workflow function called.")
    if workflow_text == "N/A" or not workflow_text.strip():
        print("No workflow data to export.")
        return None, "No workflow data to export."

    # Structure the JSON data
    workflow_data = {
        "comfy_workflow": workflow_text
    }

    # Create a temporary JSON file
    try:
        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json') as tmp_file:
            json.dump(workflow_data, tmp_file, indent=4)
            tmp_file_path = tmp_file.name
        print(f"Temporary JSON file created at: {tmp_file_path}")
    except Exception as e:
        print(f"Error creating temporary file: {e}")
        return None, "Failed to create workflow JSON file."

    # Check if the file was created successfully
    if os.path.exists(tmp_file_path):
        message = "Workflow exported successfully."
        print(message)
        return tmp_file_path, message
    else:
        message = "Failed to export workflow."
        print(message)
        return None, message

def main():
    # Create Gradio User Interface
    with gr.Blocks() as iface:
        gr.Markdown("<h1>Automatic 1111/ Comfy Metadata Reader</h1>")
        gr.Markdown("<p>Upload an image (PNG, JPEG) to extract its metadata and parse it for prompts.</p>")
        with gr.Row():
            with gr.Column():  # First column for image upload and preview
                image_input = gr.File(label="Upload Image", type="filepath")
                image_output = gr.Image(label="Image Preview")
                original_metadata_output = gr.Textbox(label="Original Metadata", lines=20, interactive=False)

            with gr.Column():  # Second column for metadata outputs
                prompt_output = gr.Textbox(label="Prompt", lines=4, show_copy_button=True, interactive=False)
                negative_prompt_output = gr.Textbox(label="Negative Prompt", lines=4, show_copy_button=True, interactive=False)
                seed_output = gr.Textbox(label="Seed Number", lines=1, show_copy_button=True, interactive=False)
                adetailer_prompt_1_output = gr.Textbox(label="ADetailer Prompt 1", lines=3, show_copy_button=True, interactive=False)
                adetailer_prompt_2_output = gr.Textbox(label="ADetailer Prompt 2", lines=3, show_copy_button=True, interactive=False)

            with gr.Column():  # Third column for additional metadata outputs
                with gr.Row():
                    comfy_prompt_output = gr.Textbox(label="Comfy Prompt", lines=4, value="N/A", interactive=False)
                with gr.Row():
                    comfy_workflow_output = gr.Textbox(label="Comfy Workflow", lines=20, value="N/A", interactive=False)
                    export_button = gr.Button("Export Workflow as JSON")
                    workflow_file = gr.File(label="Download Workflow JSON", visible=False)
                with gr.Row():
                    export_message = gr.Textbox(label="Export Status", lines=1, interactive=False, visible=False)

        # Set up the input-output relationships
        image_input.change(
            fn=extract_metadata,
            inputs=image_input,
            outputs=[
                image_output,
                prompt_output,
                negative_prompt_output,
                seed_output,
                adetailer_prompt_1_output,
                adetailer_prompt_2_output,
                original_metadata_output,
                comfy_prompt_output, 
                comfy_workflow_output,
            ]
        )

        # Connect the export button to the export function
        export_button.click(
            fn=export_workflow,
            inputs=comfy_workflow_output,
            outputs=[workflow_file, export_message],
            queue=False
        )

    iface.launch()

if __name__ == "__main__":
    main()