import gradio as gr
import piexif
import piexif.helper
import json
from PIL import Image

IGNORED_INFO_KEYS = {
    'jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',
    'loop', 'background', 'timestamp', 'duration', 'progressive', 'progression',
    'icc_profile', 'chromaticity', 'photoshop',
}

def read_info_from_image(image: Image.Image) -> tuple[str |None, dict]:
    if image is None:
        return "Please upload an image.", {}  # Return an empty dict instead of None

    items = (image.info or {}).copy()
    geninfo = items.pop('parameters', None)

    if "exif" in items:
        exif_data = items["exif"]
        try:
            exif = piexif.load(exif_data)
        except OSError:
            exif = None
        exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b'')
        try:
            exif_comment = piexif.helper.UserComment.load(exif_comment)
        except ValueError:
            exif_comment = exif_comment.decode('utf8', errors="ignore")

        if exif_comment:
            items['exif comment'] = exif_comment
            geninfo = exif_comment
    elif "comment" in items:
        geninfo = items["comment"].decode('utf8', errors="ignore")

    for field in IGNORED_INFO_KEYS:
        items.pop(field, None)

    if items.get("Software", None) == "NovelAI":
        try:
            json_info = json.loads(items["Comment"])
            sampler = "Euler a"  # Removed sd_samplers import

            geninfo = f"""{items["Description"]}
              Negative prompt: {json_info["Negative Prompt"]}
              Steps: {json_info["steps"]}, Sampler: {sampler}, CFG scale: {json_info["scale"]}, Seed: {json_info["seed"]}, Size: {image.width}x{image.height}, Clip skip: 2, ENSD: 31337"""
        except Exception as e:
            print(f"Error parsing NovelAI image generation parameters:")

    return geninfo, items

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Image Exif Parser
        [ref webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui)\n
        support png jpeg webp image format from images generated by AI tools.
        """
    )

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(sources=["upload", "clipboard"], label="Input Image", type="pil", height=680)

        with gr.Column():
            # output_metadata = gr.JSON(label="format metadata")
            output_metadata = gr.Textbox(label="format metadata")
            with gr.Accordion(open=True):
                # output_exif = gr.JSON(label="exif comments")
                output_exif = gr.Textbox(label="exif comments")

    input_image.change(
        fn=read_info_from_image,
        inputs=input_image,
        outputs=[output_metadata, output_exif],
    )

    gr.Examples(
        examples=[
            ["ex/0.png"],
            ["ex/5.jpeg"],
            ["ex/7.webp"],
            ["ex/s.png"],
        ],
        inputs=input_image,
        outputs=[output_metadata, output_exif],
        fn=read_info_from_image,
        cache_examples=False,
        label="Exmaple format: png, jpeg, webp"
    )

demo.launch()