"""保存图像节点 - 支持 PNG/JPEG/WebP 格式输出"""

import os
import json
import numpy as np
from PIL import Image
from PIL.PngImagePlugin import PngInfo

import folder_paths
from comfy.cli_args import args


class SaveImageFormat:
    """保存图像，支持 PNG / JPEG / WebP 三种格式"""

    FORMATS = ["PNG", "JPEG", "WebP"]

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

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "图像": ("IMAGE",),
                "文件名前缀": ("STRING", {"default": "ComfyUI"}),
                "输出格式": (cls.FORMATS, {"default": "PNG"}),
                "质量": ("INT", {
                    "default": 100, "min": 1, "max": 100, "step": 1,
                    "tooltip": "图片质量。100=不压缩（JPEG 最高质量 / WebP 无损）；"
                               "小于 100 时按该数值压缩（如 90），仅对 JPEG / WebP 生效。",
                }),
            },
            "optional": {
                "保存路径": ("STRING", {"default": ""}),
            },
            "hidden": {
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO",
            },
        }

    RETURN_TYPES = ()
    FUNCTION = "save_images"
    OUTPUT_NODE = True
    CATEGORY = "image"
    DESCRIPTION = "保存图像，支持 PNG / JPEG / WebP 格式输出。"

    _EXT_MAP = {"PNG": ".png", "JPEG": ".jpg", "WebP": ".webp"}

    @classmethod
    def IS_CHANGED(cls, 图像=None, **kwargs):
        """保存节点属于有副作用的输出节点，不能复用上次的缓存结果。"""
        return float("nan")

    def save_images(self, 图像=None, 文件名前缀="ComfyUI", 输出格式="PNG",
                    质量=100, 保存路径="", prompt=None, extra_pnginfo=None):
        images = 图像
        filename_prefix = 文件名前缀
        format = 输出格式
        quality = int(质量)
        custom_dir = (保存路径 or "").strip()

        if custom_dir:
            # 保存到用户指定的文件夹。自定义路径不会经过
            # folder_paths.get_save_image_path()，因此需要在此处自行避让重名。
            full_output_folder = custom_dir
            os.makedirs(full_output_folder, exist_ok=True)
            filename = filename_prefix
            counter = 1
            subfolder = ""
        else:
            full_output_folder, filename, counter, subfolder, filename_prefix = \
                folder_paths.get_save_image_path(
                    filename_prefix, self.output_dir,
                    images[0].shape[1], images[0].shape[0]
                )

        ext = self._EXT_MAP.get(format, ".png")
        results = []

        # 为整个输入批次预留一段连续编号，避免重复运行时覆盖已有文件。
        # 默认 output 路径和自定义保存路径都在这里复核一次；这样即使外部
        # 编号器返回了已使用的计数，也不会覆盖。兼容 %batch_num% 占位符。
        while any(
            os.path.exists(
                os.path.join(
                    full_output_folder,
                    f"{filename.replace('%batch_num%', str(batch_number))}_{counter + batch_number:05}_{ext}",
                )
            )
            for batch_number in range(len(images))
        ):
            counter += 1

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

            filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
            file = f"{filename_with_batch_num}_{counter:05}_{ext}"

            filepath = os.path.join(full_output_folder, file)

            if format == "PNG":
                metadata = None
                if not args.disable_metadata:
                    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]))
                img.save(filepath, pnginfo=metadata,
                         compress_level=self.compress_level)
            elif format == "JPEG":
                if img.mode == "RGBA":
                    img = img.convert("RGB")
                if quality >= 100:
                    img.save(filepath, quality=100, optimize=True)
                else:
                    img.save(filepath, quality=quality, optimize=True)
            elif format == "WebP":
                if quality >= 100:
                    img.save(filepath, lossless=True)
                else:
                    img.save(filepath, quality=quality, method=6)

            results.append({
                "filename": file,
                "subfolder": subfolder,
                "type": self.type,
            })
            counter += 1

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