import os
import json
import re
import numpy as np
from datetime import datetime

import folder_paths
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from comfy.cli_args import args


def _create_webp_metadata(img, prompt, extra_pnginfo):
    """Creates EXIF metadata for WebP images. Reference: WAS_Node_Suite.py:7452-7464"""
    exif_data = img.getexif()
    if prompt is not None:
        exif_data[0x010f] = "Prompt:" + json.dumps(prompt)  # EXIF 0x010f = Make
    if extra_pnginfo is not None:
        workflow_metadata = ''
        for key, value in extra_pnginfo.items():
            workflow_metadata += json.dumps(value)
        exif_data[0x010e] = "Workflow:" + workflow_metadata  # EXIF 0x010e = ImageDescription
    return exif_data.tobytes()


class SaveImageAsSDWebUI:
    def __init__(self):
        self.output_dir = folder_paths.get_output_directory()
        self.type = "output"
        self.compress_level = 4

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE", {"tooltip": "The images to save."}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Seed value used in the filename."}),
                "filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "Prefix for the filename. Can include path separators (e.g. 'comfy-txt2img/') for subdirectory classification."}),
                "image_format": (["png", "jpg", "webp"], {"default": "png", "tooltip": "Image save format."}),
                "quality": ("INT", {"default": 100, "min": 0, "max": 100, "tooltip": "Save quality for JPG/WebP formats (0-100). Ignored for PNG."}),
            },
            "hidden": {
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO",
            },
        }

    RETURN_TYPES = ()
    FUNCTION = "save_images"
    OUTPUT_NODE = True
    CATEGORY = "image"

    def save_images(self, images, seed, filename_prefix="ComfyUI", image_format="png", quality=100, prompt=None, extra_pnginfo=None):
        # Split prefix: "dir/name" -> dir="dir", name="name"
        # "dir/" -> dir="dir", name=""
        # "name" -> dir="", name="name"
        prefix = filename_prefix
        if '/' in prefix or '\\' in prefix:
            prefix = prefix.replace('\\', '/')
            last_sep = prefix.rfind('/')
            folder_parts = prefix[:last_sep].split('/')
            name_part = prefix[last_sep + 1:]
        else:
            folder_parts = []
            name_part = prefix

        date_str = datetime.now().strftime("%Y-%m-%d")
        output_folder = os.path.join(self.output_dir, *folder_parts, date_str)
        os.makedirs(output_folder, exist_ok=True)

        # Safety check: prevent path traversal
        if os.path.commonpath((self.output_dir, os.path.abspath(output_folder))) != self.output_dir:
            raise Exception("Saving outside the output folder is not allowed.")

        # Scan for max counter
        if name_part:
            pattern = re.compile(r'^' + re.escape(name_part) + r'(\d{5})-\d+\.\w+$')
        else:
            pattern = re.compile(r'^(\d{5})-\d+\.\w+$')
        max_counter = 0
        for f in os.listdir(output_folder):
            m = pattern.match(f)
            if m:
                max_counter = max(max_counter, int(m.group(1)))
        counter = max_counter + 1

        subfolder = os.path.join(*folder_parts, date_str) if folder_parts else date_str

        results = []
        for (batch_number, image) in enumerate(images):
            i = 255. * image.cpu().numpy()
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))

            metadata = None
            if not args.disable_metadata:
                if image_format == "png":
                    metadata = PngInfo()
                    if prompt is not None:
                        metadata.add_text("prompt", json.dumps(prompt))
                    if extra_pnginfo is not None:
                        for x in extra_pnginfo:
                            metadata.add_text(x, json.dumps(extra_pnginfo[x]))
                elif image_format == "webp":
                    metadata = _create_webp_metadata(img, prompt, extra_pnginfo)

            file = f"{name_part}{counter:05d}-{seed}.{image_format}"
            filepath = os.path.join(output_folder, file)

            save_kwargs = {}
            if image_format == "png":
                save_kwargs["pnginfo"] = metadata
                save_kwargs["compress_level"] = self.compress_level
            elif image_format == "webp":
                save_kwargs["exif"] = metadata
                save_kwargs["quality"] = quality
                if img.mode == "RGBA":
                    img = img.convert("RGB")
            elif image_format == "jpg":
                save_kwargs["quality"] = quality
                if img.mode == "RGBA":
                    img = img.convert("RGB")

            img.save(filepath, **save_kwargs)
            results.append({
                "filename": file,
                "subfolder": subfolder,
                "type": self.type,
            })
            counter += 1

        return {"ui": {"images": results}}


NODE_CLASS_MAPPINGS = {
    "SaveImageAsSDWebUI": SaveImageAsSDWebUI,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "SaveImageAsSDWebUI": "Save Image (SD WebUI Style)",
}