import gradio as gr
from PIL import Image
import json

def extract_parameters(image, show_full_metadata=False):
    """
    Extract and format parameters from a ComfyUI image.
    If show_full_metadata is False, shows only the main prompt cleanly.
    If True, shows the full metadata with better formatting.
    """
    try:
        # Open the uploaded image
        img = Image.open(image)
        # Get metadata from the image's info dictionary
        metadata = img.info

        if not metadata:
            return "No parameters found in image metadata."

        # Check if "prompt" key exists (common in ComfyUI)
        if "prompt" in metadata:
            raw_prompt = metadata["prompt"]
            try:
                # Parse as JSON if possible
                prompt_data = json.loads(raw_prompt)
                # Look for the main prompt text (often in a node like '51')
                main_prompt = None
                for node_id, node in prompt_data.items():
                    if "inputs" in node and "text" in node["inputs"]:
                        # Prioritize nodes with meaningful text input
                        text = node["inputs"]["text"]
                        if isinstance(text, str) and text.strip():
                            main_prompt = text
                            break
                
                if not main_prompt:
                    main_prompt = "No clear prompt found in workflow."

                if show_full_metadata:
                    # Format full metadata with spacing
                    full_output = "Full Workflow Metadata:\n\n"
                    for node_id, node in prompt_data.items():
                        full_output += f"Node {node_id}:\n"
                        for key, value in node.items():
                            full_output += f"  {key}: {value}\n"
                        full_output += "\n"
                    return f"Main Prompt:\n{main_prompt}\n\n{full_output}"
                else:
                    # Just the cleaned-up main prompt
                    return f"Main Prompt:\n{main_prompt}"
            except json.JSONDecodeError:
                # If not JSON, treat as raw text
                if show_full_metadata:
                    return f"Raw Prompt:\n{raw_prompt}\n\nFull Metadata:\n{metadata}"
                return f"Main Prompt:\n{raw_prompt}"
        else:
            # No "prompt" key, show all metadata if toggled
            if show_full_metadata:
                output = "Full Metadata:\n\n"
                for key, value in metadata.items():
                    output += f"{key}:\n  {value}\n\n"
                return output
            return "No prompt found, toggle 'Show Full Metadata' for details."

    except Exception as e:
        return f"Error processing image: {str(e)}"

# Define the Gradio interface
interface = gr.Interface(
    fn=extract_parameters,
    inputs=[
        gr.Image(type="filepath", label="Upload a ComfyUI Image"),
        gr.Checkbox(label="Show Full Metadata", value=False)  # Changed 'default' to 'value'
    ],
    outputs=gr.Textbox(label="Extracted Parameters", lines=10),
    title="ComfyUI Image to Parameters",
    description="Upload a ComfyUI image to extract its prompt. Toggle 'Show Full Metadata' for the complete workflow.",
    theme="default"
)

# Launch the app
interface.launch()