import gradio as gr
from PIL import Image, ImageFilter
from PIL.ExifTags import TAGS
import numpy as np
import re


def main():
    with gr.Blocks() as iface:
        gr.Markdown("# Image Metadata Extractor")
        gr.Markdown("Upload an image (PNG, JPEG) to extract its metadata and parse it for prompts.\n")
        
        with gr.Row():
            with gr.Column():
                image_input = gr.File(label="Upload Image", type="filepath")
                censor_button = gr.Button("👁️ Toggle Censor")
                image_output = gr.Image(label="Image Preview")
                original_metadata_output = gr.Textbox(label="Original Metadata", lines=20)

            with gr.Column():
                prompt_output = gr.Textbox(label="Prompt", lines=4, show_copy_button=True)
                negative_prompt_output = gr.Textbox(label="Negative Prompt", lines=4, show_copy_button=True)
                seed_output = gr.Textbox(label="Seed Number", lines=1, show_copy_button=True)
                adetailer_prompt_1_output = gr.Textbox(label="ADetailer Prompt 1", lines=3, show_copy_button=True)
                adetailer_prompt_2_output = gr.Textbox(label="ADetailer Prompt 2", lines=3, show_copy_button=True)
                adetailer_negative_prompt_output = gr.Textbox(label="ADetailer Negative Prompt", lines=4, show_copy_button=True)
                original_prompt_output = gr.Textbox(label="Original Prompt", lines=4, show_copy_button=True)
                hires_prompt_output = gr.Textbox(label="Hires Prompt", lines=4, show_copy_button=True)

        # State to store original and censored images
        image_state = gr.State(value={"original": None, "censored": None, "current": None})

        # Image upload handler
        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,
                adetailer_negative_prompt_output,
                original_prompt_output,
                hires_prompt_output,
                original_metadata_output,
                image_state,
            ],
        )

        # Button to toggle censor
        censor_button.click(
            fn=toggle_censor,
            inputs=image_state,
            outputs=image_output,
        )

    iface.launch(show_error=True)


def extract_metadata(image_file):
    img = Image.open(image_file)
    metadata = {}

    if img.format == "PNG":
        metadata = img.info
    elif img.format in ["JPEG"]:
        exif_data = img._getexif()
        if exif_data:
            for tag, value in exif_data.items():
                tag_name = TAGS.get(tag, tag)
                metadata[tag_name] = value

    # Generate censored version of the image
    censored_img = img.filter(ImageFilter.GaussianBlur(40))
    img_display = np.array(img)

    prompt = "N/A"
    negative_prompt = "N/A"
    adetailer_prompt_1 = "N/A"
    adetailer_prompt_2 = "N/A"
    adetailer_negative_prompt = "N/A"
    original_prompt = "N/A"
    hires_prompt = "N/A"
    seed_number = -1

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

    if "parameters" in metadata:
        parameters = metadata["parameters"]

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

        negative_prompt_match = re.search(r"Negative prompt:(.*?)(Steps:|$)", parameters, re.DOTALL)
        if negative_prompt_match:
            negative_prompt = negative_prompt_match.group(1).strip()

        adetailer_prompt_1_match = re.search(r'ADetailer prompt:\s*"(.*?)"', parameters)
        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)
        if adetailer_prompt_2_match:
            adetailer_prompt_2 = adetailer_prompt_2_match.group(1).strip()

        adetailer_negative_prompt_match = re.search(r'ADetailer negative prompt:\s*"(.*?)"', parameters)
        if adetailer_negative_prompt_match:
            adetailer_negative_prompt = adetailer_negative_prompt_match.group(1).strip()

        original_prompt_match = re.search(r'DTG prompt:\s*"(.*?)"', parameters, re.DOTALL)
        if original_prompt_match:
            original_prompt = original_prompt_match.group(1).strip()

        hires_prompt_match = re.search(r'Hires prompt:\s*"(.*?)"', parameters, re.DOTALL)
        if hires_prompt_match:
            hires_prompt = hires_prompt_match.group(1).strip()

        seed_match = re.search(r"Seed:\s*(\d+)", parameters)
        if seed_match:
            seed_number = seed_match.group(1).strip()

    # Return data along with both original and censored images
    return (
        img_display,
        prompt,
        negative_prompt,
        seed_number,
        adetailer_prompt_1,
        adetailer_prompt_2,
        adetailer_negative_prompt,
        original_prompt,
        hires_prompt,
        metadata_str,
        {"original": img, "censored": censored_img, "current": img},
    )


def toggle_censor(image_state):
    """
    Toggle between the original and censored image.
    """
    if image_state["current"] == image_state["original"]:
        image_state["current"] = image_state["censored"]
    else:
        image_state["current"] = image_state["original"]
    
    return np.array(image_state["current"])


if __name__ == "__main__":
    main()