clean_save_node.py

"""ComfyUI custom node: save images with enforced metadata sanitization."""

from __future__ import annotations

import base64
import io
import os
import re
import time
import zlib
from typing import Any, Dict, List, Optional, Tuple

import random

import numpy as np
from PIL import Image
from .metadata.randomize import build_batch_context, resolve_settings
from .metadata.exif_builder import build_exif
from .metadata.xmp_builder import build_xmp
from .metadata.iptc_builder import build_iptc

try:
    import torch
except Exception:  # pragma: no cover - torch is always available in ComfyUI
    torch = None  # type: ignore

try:
    import folder_paths
except Exception:  # pragma: no cover - fallback for local execution
    folder_paths = None  # type: ignore


PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
PNG_TEXT_METADATA_CHUNKS = {b"tEXt", b"zTXt", b"iTXt"}


def _tensor_to_pil(image_tensor: Any) -> Image.Image:
    if torch is None:
        raise RuntimeError("Torch not available; cannot convert ComfyUI IMAGE tensor.")

    if isinstance(image_tensor, torch.Tensor):
        image = image_tensor.detach().cpu().numpy()
    else:
        image = np.asarray(image_tensor)

    image = np.clip(image * 255.0, 0, 255).astype(np.uint8)
    if image.shape[-1] == 4:
        return Image.fromarray(image, mode="RGBA")
    return Image.fromarray(image, mode="RGB")


def _print_debug_console(
    filename: str,
    mode: str,
    exif_bytes: Optional[bytes],
    xmp_bytes: Optional[bytes],
    iptc_iim: Optional[bytes],
) -> None:
    has_exif = "yes" if exif_bytes else "no"
    has_xmp = "yes" if xmp_bytes else "no"
    has_iptc = "yes" if iptc_iim else "no"
    print(
        f"SaveImageCleanMetadata[debug]: file={filename} mode={mode} "
        f"exif={has_exif} xmp={has_xmp} iptc={has_iptc}"
    )


def _png_chunk(chunk_type: bytes, data: bytes) -> bytes:
    length = len(data).to_bytes(4, "big")
    crc = zlib.crc32(chunk_type + data) & 0xFFFFFFFF
    return length + chunk_type + data + crc.to_bytes(4, "big")


def _build_itxt_chunk(keyword: str, text_bytes: bytes) -> bytes:
    keyword_bytes = keyword.encode("latin-1")[:79]
    data = (
        keyword_bytes
        + b"\x00"
        + b"\x00"
        + b"\x00"
        + b"\x00"
        + b"\x00"
        + text_bytes
    )
    return _png_chunk(b"iTXt", data)


def _process_png_bytes(
    png_bytes: bytes,
    xmp_bytes: Optional[bytes],
    iptc_iim: Optional[bytes],
    remove_exif: bool = True,
) -> bytes:
    if not png_bytes.startswith(PNG_SIGNATURE):
        return png_bytes

    xmp_chunk = None
    iptc_chunk = None
    if xmp_bytes:
        xmp_chunk = _build_itxt_chunk("XML:com.adobe.xmp", xmp_bytes)
    if iptc_iim:
        iptc_chunk = _build_itxt_chunk("IPTC", base64.b64encode(iptc_iim))

    new_chunks: List[bytes] = []
    inserted = False

    idx = len(PNG_SIGNATURE)
    length_bytes = len(png_bytes)
    while idx + 12 <= length_bytes:
        chunk_length = int.from_bytes(png_bytes[idx : idx + 4], "big")
        chunk_type = png_bytes[idx + 4 : idx + 8]
        chunk_data_end = idx + 8 + chunk_length
        chunk_end = chunk_data_end + 4
        if chunk_end > length_bytes:
            break

        if chunk_type in PNG_TEXT_METADATA_CHUNKS:
            idx = chunk_end
            continue

        if remove_exif and chunk_type == b"eXIf":
            idx = chunk_end
            continue

        if not inserted and chunk_type == b"IDAT":
            if xmp_chunk:
                new_chunks.append(xmp_chunk)
            if iptc_chunk:
                new_chunks.append(iptc_chunk)
            inserted = True

        if not inserted and chunk_type == b"IEND":
            if xmp_chunk:
                new_chunks.append(xmp_chunk)
            if iptc_chunk:
                new_chunks.append(iptc_chunk)
            inserted = True

        new_chunks.append(png_bytes[idx:chunk_end])
        idx = chunk_end

    if not inserted:
        if xmp_chunk:
            new_chunks.append(xmp_chunk)
        if iptc_chunk:
            new_chunks.append(iptc_chunk)

    return PNG_SIGNATURE + b"".join(new_chunks)


def _add_png_title_subject(
    png_bytes: bytes,
    title: str,
    subject: str = "",
) -> bytes:
    if not png_bytes.startswith(PNG_SIGNATURE):
        return png_bytes

    title_chunk = _build_itxt_chunk("Title", title.encode("utf-8"))
    subject_chunk = _build_itxt_chunk("Subject", subject.encode("utf-8"))

    new_chunks: List[bytes] = []
    inserted = False

    idx = len(PNG_SIGNATURE)
    length_bytes = len(png_bytes)
    while idx + 12 <= length_bytes:
        chunk_length = int.from_bytes(png_bytes[idx : idx + 4], "big")
        chunk_type = png_bytes[idx + 4 : idx + 8]
        chunk_data_end = idx + 8 + chunk_length
        chunk_end = chunk_data_end + 4
        if chunk_end > length_bytes:
            break

        if not inserted and chunk_type == b"IDAT":
            new_chunks.append(title_chunk)
            new_chunks.append(subject_chunk)
            inserted = True

        if not inserted and chunk_type == b"IEND":
            new_chunks.append(title_chunk)
            new_chunks.append(subject_chunk)
            inserted = True

        new_chunks.append(png_bytes[idx:chunk_end])
        idx = chunk_end

    if not inserted:
        new_chunks.append(title_chunk)
        new_chunks.append(subject_chunk)

    return PNG_SIGNATURE + b"".join(new_chunks)


def _encode_png_bytes(img: Image.Image, exif_bytes: Optional[bytes] = None) -> bytes:
    buffer = io.BytesIO()
    save_kwargs: Dict[str, Any] = {
        "format": "PNG",
        "compress_level": 9,
    }
    if exif_bytes:
        save_kwargs["exif"] = exif_bytes

    try:
        img.save(buffer, **save_kwargs)
    except TypeError:
        if exif_bytes:
            save_kwargs.pop("exif", None)
            buffer = io.BytesIO()
            img.save(buffer, **save_kwargs)
            print(
                "SaveImageCleanMetadata: PNG EXIF chunk not supported by this Pillow build; "
                "saved without EXIF."
            )
        else:
            raise

    return buffer.getvalue()


def _get_output_dir() -> str:
    if folder_paths is not None:
        return folder_paths.get_output_directory()
    return os.path.join(os.getcwd(), "output")


def _get_save_image_path(
    filename_prefix: str,
    output_dir: str,
    image_width: int,
    image_height: int,
) -> Tuple[str, str, int, str, str]:
    if folder_paths is not None:
        return folder_paths.get_save_image_path(
            filename_prefix,
            output_dir,
            image_width,
            image_height,
        )

    # Local fallback that mirrors ComfyUI's path/counter behavior.
    if "%" in filename_prefix:
        now = time.localtime()
        filename_prefix = filename_prefix.replace("%width%", str(image_width))
        filename_prefix = filename_prefix.replace("%height%", str(image_height))
        filename_prefix = filename_prefix.replace("%year%", str(now.tm_year))
        filename_prefix = filename_prefix.replace("%month%", str(now.tm_mon).zfill(2))
        filename_prefix = filename_prefix.replace("%day%", str(now.tm_mday).zfill(2))
        filename_prefix = filename_prefix.replace("%hour%", str(now.tm_hour).zfill(2))
        filename_prefix = filename_prefix.replace("%minute%", str(now.tm_min).zfill(2))
        filename_prefix = filename_prefix.replace("%second%", str(now.tm_sec).zfill(2))

    normalized = os.path.normpath(filename_prefix)
    subfolder = os.path.dirname(normalized)
    filename = os.path.basename(normalized)

    output_abs = os.path.abspath(output_dir)
    full_output_folder = os.path.join(output_abs, subfolder)
    full_output_abs = os.path.abspath(full_output_folder)

    if os.path.commonpath((output_abs, full_output_abs)) != output_abs:
        raise Exception("Saving image outside the output folder is not allowed.")

    os.makedirs(full_output_abs, exist_ok=True)

    prefix_len = len(os.path.basename(filename_prefix))
    counter = 1
    for entry in os.listdir(full_output_abs):
        prefix = entry[: prefix_len + 1]
        try:
            digits = int(entry[prefix_len + 1 :].split("_")[0])
        except Exception:
            continue
        if os.path.normcase(prefix[:-1]) == os.path.normcase(filename) and prefix[-1] == "_":
            counter = max(counter, digits + 1)

    return full_output_abs, filename, counter, subfolder, filename_prefix


def _batch_images(images: Any) -> Any:
    if torch is not None and isinstance(images, torch.Tensor):
        return images
    if torch is not None:
        items = list(images)
        if items and isinstance(items[0], torch.Tensor):
            return torch.stack(items)
        return np.stack(items)
    return np.stack(list(images))


def _normalize_filename_prefix(value: Any) -> str:
    if not isinstance(value, str):
        return "ComfyUI"
    normalized = value.strip()
    if not normalized:
        return "ComfyUI"
    if normalized.lower() in {"false", "none", "null", "no"}:
        return "ComfyUI"
    return normalized


def _camera_default_path(
    output_dir: str, filename_prefix: str, profile, image_width: int, image_height: int
):
    full_output_folder, _, _, _, _ = _get_save_image_path(
        filename_prefix=filename_prefix,
        output_dir=output_dir,
        image_width=image_width,
        image_height=image_height,
    )
    prefix = profile.default_filename_prefix
    digits = profile.default_filename_counter_digits
    counter = 1
    pattern_re = re.compile(rf"^{re.escape(prefix)}(\d{{{digits}}})_\.png$")
    os.makedirs(full_output_folder, exist_ok=True)
    for entry in os.listdir(full_output_folder):
        m = pattern_re.match(entry)
        if m:
            counter = max(counter, int(m.group(1)) + 1)
    return full_output_folder, prefix, counter


class SaveImageCleanMetadata:
    """Save ComfyUI images to disk with enforced metadata sanitization."""

    @classmethod
    def INPUT_TYPES(cls) -> Dict[str, Dict[str, Any]]:
        return {
            "required": {
                "images": ("IMAGE",),
                "filename_prefix": ("STRING", {"default": "ComfyUI"}),
                "use_camera_default_filename": ("BOOLEAN", {"default": False}),
                "strip_all_metadata": ("BOOLEAN", {"default": False}),
                "debug_console": ("BOOLEAN", {"default": False}),
            },
            "optional": {
                "metadata_settings": ("METADATA_SETTINGS",),
            },
        }

    RETURN_TYPES: Tuple[()] = ()
    FUNCTION = "save_images"
    OUTPUT_NODE = True
    CATEGORY = "image"

    def save_images(
        self,
        images: Any,
        filename_prefix: str,
        use_camera_default_filename: bool,
        strip_all_metadata: bool,
        debug_console: bool,
        metadata_settings: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        filename_prefix = _normalize_filename_prefix(filename_prefix)

        if use_camera_default_filename and (metadata_settings is None or strip_all_metadata):
            raise RuntimeError(
                "Cannot use camera default filename without a ControlMetadataSettings node "
                "or when strip_all_metadata is enabled"
            )

        apply_metadata = (not strip_all_metadata) and (metadata_settings is not None)

        output_dir = _get_output_dir()
        os.makedirs(output_dir, exist_ok=True)

        batch = _batch_images(images)

        full_output_folder, filename, counter, subfolder, _ = _get_save_image_path(
            filename_prefix=filename_prefix,
            output_dir=output_dir,
            image_width=batch[0].shape[1],
            image_height=batch[0].shape[0],
        )

        batch_context = None
        rng = None
        if apply_metadata:
            seed_raw = int(metadata_settings.get("seed", 0))
            rng_seed = seed_raw if seed_raw != 0 else int.from_bytes(os.urandom(8), "big")
            rng = random.Random(rng_seed)
            batch_context = build_batch_context(
                metadata_settings, rng, batch_size=int(batch.shape[0])
            )

        if use_camera_default_filename:
            profile = batch_context["profile"]
            full_output_folder, filename, counter = _camera_default_path(
                output_dir, filename_prefix, profile,
                image_width=batch[0].shape[1],
                image_height=batch[0].shape[0],
            )

        results: List[Dict[str, str]] = []

        for i in range(batch.shape[0]):
            filename_with_batch_num = filename.replace("%batch_num%", str(i))
            if use_camera_default_filename:
                profile = batch_context["profile"]
                digits = profile.default_filename_counter_digits
                file = f"{filename_with_batch_num}{counter:0{digits}}_.png"
            else:
                file = f"{filename_with_batch_num}_{counter:05}_.png"
            title_value = file

            img = _tensor_to_pil(batch[i])
            img.info.clear()

            exif_bytes = None
            xmp_bytes = None
            iptc_bytes = None
            if apply_metadata:
                resolved = resolve_settings(metadata_settings, rng, i, batch_context)
                exif_bytes = build_exif(resolved)
                xmp_bytes = build_xmp(resolved)
                iptc_bytes = build_iptc(resolved)

            image_bytes = _encode_png_bytes(img, exif_bytes=exif_bytes)
            image_bytes = _process_png_bytes(
                image_bytes, xmp_bytes, iptc_bytes, remove_exif=(not apply_metadata)
            )
            image_bytes = _add_png_title_subject(
                image_bytes, title=title_value, subject=""
            )

            filepath = os.path.join(full_output_folder, file)
            with open(filepath, "wb") as outfile:
                outfile.write(image_bytes)

            if debug_console:
                mode = "strip" if not apply_metadata else "apply"
                _print_debug_console(
                    filename=os.path.join(subfolder, file) if subfolder else file,
                    mode=mode,
                    exif_bytes=exif_bytes,
                    xmp_bytes=xmp_bytes,
                    iptc_iim=iptc_bytes,
                )

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

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


NODE_CLASS_MAPPINGS = {
    "SaveImageCleanMetadata": SaveImageCleanMetadata,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SaveImageCleanMetadata": "Save Image (Control Metadata)",
}