from inspect import cleandoc
import os
import glob

try:
    from comfy.comfy_types.node_typing import IO, ComfyNodeABC
except:

    class IO:
        BOOLEAN = "BOOLEAN"
        INT = "INT"
        FLOAT = "FLOAT"
        STRING = "STRING"
        NUMBER = "FLOAT,INT"
        IMAGE = "IMAGE"
        MASK = "MASK"
        ANY = "*"

    ComfyNodeABC = object

try:
    from folder_paths import get_input_directory, get_output_directory
except:

    def get_input_directory():
        return "./"

    get_output_directory = get_input_directory

# ComfyUI's official prefix->path resolver is reused when running under ComfyUI
# so the nodes follow upstream naming exactly (folder layout, %width%/%height%/
# date token expansion, auto-incrementing counters and containment checks).
# Outside ComfyUI (e.g. the standalone test-suite) `folder_paths` is not
# importable, so a minimal local equivalent is used as a fallback only.
try:
    from folder_paths import get_save_image_path as _official_get_save_image_path
except Exception:
    _official_get_save_image_path = None


def _require_numpy():
    try:
        import numpy as np

        return np
    except ModuleNotFoundError as e:
        raise ModuleNotFoundError(
            "basic_data_handling: Missing dependency 'numpy'. It seems your ComfyUI installation is faulty."
            "Only for development purposes: Install it with `pip install .[dev] numpy torch pillow` or `pip install numpy`."
        ) from e


def _require_torch():
    try:
        import torch

        return torch
    except ModuleNotFoundError as e:
        raise ModuleNotFoundError(
            "basic_data_handling: Missing dependency 'torch'. It seems your ComfyUI installation is faulty."
            "Only for development purposes: Install it with `pip install .[dev] numpy torch pillow` or `pip install torch`."
        ) from e


def _require_pillow():
    try:
        from PIL import Image, ImageOps

        return Image, ImageOps
    except ModuleNotFoundError as e:
        raise ModuleNotFoundError(
            "basic_data_handling: Missing dependency 'pillow'. It seems your ComfyUI installation is faulty."
            "Only for development purposes: Install it with `pip install .[dev] numpy torch pillow` or `pip install pillow`."
        ) from e


# helper functions:


def load_image_helper(path: str):
    """Helper function to load an image from a path"""
    Image, ImageOps = _require_pillow()
    try:
        import pillow_jxl  # noqa: F401 - imported but unused, kept for JPEG XL support
    except ModuleNotFoundError:
        pass

    if not os.path.exists(path):
        return None

    # Open and process the image
    try:
        img = Image.open(path)
        img = ImageOps.exif_transpose(img)
        return img
    except Exception:
        return None


def extract_mask_from_alpha(img):
    """Extract a mask from the alpha channel of an image"""
    np = _require_numpy()
    torch = _require_torch()

    if "A" in img.getbands():
        alpha = np.array(img.getchannel("A")).astype(np.float32) / 255.0
        mask_tensor = 1.0 - torch.from_numpy(alpha)
    elif img.mode == "P" and "transparency" in img.info:
        alpha = np.array(img.convert("RGBA").getchannel("A")).astype(np.float32) / 255.0
        mask_tensor = 1.0 - torch.from_numpy(alpha)
    else:
        # Create a blank mask if no alpha channel
        mask_tensor = torch.zeros((img.height, img.width), dtype=torch.float32)

    # Add batch dimension to mask
    mask_tensor = mask_tensor.unsqueeze(0)

    return mask_tensor


def extract_mask_from_greyscale(img):
    """Extract a mask from a greyscale image or the red channel of an RGB image"""
    np = _require_numpy()
    torch = _require_torch()

    if img.mode == "L":
        # Image is already greyscale
        gray = np.array(img).astype(np.float32) / 255.0
    elif img.mode == "RGB" or img.mode == "RGBA":
        # Use the red channel of RGB or RGBA
        gray = np.array(img.getchannel("R")).astype(np.float32) / 255.0
    else:
        # Convert to greyscale if it's another format
        gray_img = img.convert("L")
        gray = np.array(gray_img).astype(np.float32) / 255.0

    # Convert to tensor and invert (white pixels in image = transparent in mask)
    mask_tensor = 1.0 - torch.from_numpy(gray)

    # Add batch dimension
    mask_tensor = mask_tensor.unsqueeze(0)

    return mask_tensor


# the nodes:


class PathAbspath(ComfyNodeABC):
    """
    Returns the absolute path of a file or directory.

    This node takes a path and returns its absolute (full) path
    by resolving any relative path components and symbolic links.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path to resolve to an absolute path."}),
            }
        }

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("absolute path",)
    OUTPUT_TOOLTIPS = ("The absolute path with relative components and symlinks resolved.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "get_abspath"

    def get_abspath(self, path: str) -> tuple[str]:
        return (os.path.abspath(path),)


class PathBasename(ComfyNodeABC):
    """
    Returns the base name of a path.

    This node extracts the filename component from a path,
    removing any directory information.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path whose final filename component is returned."}),
            }
        }

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("basename",)
    OUTPUT_TOOLTIPS = ("The final filename component of the path.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "get_basename"

    def get_basename(self, path: str) -> tuple[str]:
        return (os.path.basename(path),)


class PathCommonPrefix(ComfyNodeABC):
    """
    Finds the common prefix of multiple paths.

    This node returns the longest common leading component of the given paths.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path1": (IO.STRING, {"default": "", "tooltip": "First path."}),
            },
            "optional": {
                "path2": (IO.STRING, {"default": "", "tooltip": "Second path."}),
            },
        }

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("common prefix",)
    OUTPUT_TOOLTIPS = ("The longest common leading component of the given paths.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "get_common_prefix"

    def get_common_prefix(self, path1: str, path2: str = "") -> tuple[str]:
        paths = [p for p in [path1, path2] if p]
        return (os.path.commonprefix(paths),)


class PathDirname(ComfyNodeABC):
    """
    Returns the directory name of a path.

    This node extracts the directory component from a path,
    removing the filename.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path whose directory component is returned."}),
            }
        }

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("dirname",)
    OUTPUT_TOOLTIPS = ("The directory (parent) component of the path.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "get_dirname"

    def get_dirname(self, path: str) -> tuple[str]:
        return (os.path.dirname(path),)


class PathExists(ComfyNodeABC):
    """
    Checks if a path exists in the filesystem.

    This node returns True if the path exists (either as a file or a directory),
    and False otherwise.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path to check for existence."}),
            }
        }

    RETURN_TYPES = (IO.BOOLEAN,)
    RETURN_NAMES = ("exists",)
    OUTPUT_TOOLTIPS = ("True when the path exists as a file or directory.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "check_exists"

    def check_exists(self, path: str) -> tuple[bool]:
        return (os.path.exists(path),)


class PathExpandVars(ComfyNodeABC):
    """
    Expands environment variables in a path.

    This node replaces environment variables in a path with their values.
    For example, $HOME or ${HOME} on Unix, or %USERPROFILE% on Windows.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path that may contain environment variables (e.g. $HOME, %USERPROFILE%)."}),
            }
        }

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("expanded path",)
    OUTPUT_TOOLTIPS = ("The path with environment variables expanded.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "expand_vars"

    def expand_vars(self, path: str) -> tuple[str]:
        return (os.path.expandvars(path),)


class PathGetCwd(ComfyNodeABC):
    """
    Returns the current working directory.

    This node returns the current working directory as an absolute path.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {}}

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("current directory",)
    OUTPUT_TOOLTIPS = ("The current working directory as an absolute path.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "get_cwd"

    def get_cwd(self) -> tuple[str]:
        return (os.getcwd(),)


class PathGetExtension(ComfyNodeABC):
    """
    Returns the extension of a file.

    This node extracts the file extension from a path,
    including the dot (e.g., '.txt').
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path whose file extension is extracted."}),
            }
        }

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("extension",)
    OUTPUT_TOOLTIPS = ("The extension including the dot (e.g. '.txt'); empty when there is none.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "get_extension"

    def get_extension(self, path: str) -> tuple[str]:
        return (os.path.splitext(path)[1],)


class PathGetSize(ComfyNodeABC):
    """
    Returns the size of a file in bytes.

    This node returns the size in bytes of the file at the given path.
    Raises an error if the path doesn't exist or isn't a file.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "File whose size in bytes is returned."}),
            }
        }

    RETURN_TYPES = (IO.INT,)
    RETURN_NAMES = ("size (bytes)",)
    OUTPUT_TOOLTIPS = ("The file size in bytes.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "get_size"

    def get_size(self, path: str) -> tuple[int]:
        if not os.path.exists(path):
            raise FileNotFoundError(f"Basic data handling: Path does not exist: {path}")
        if not os.path.isfile(path):
            raise ValueError(f"Basic data handling: Path is not a file: {path}")
        return (os.path.getsize(path),)


class PathGlob(ComfyNodeABC):
    """
    Finds paths matching a pattern.

    This node returns a list of paths matching the given pattern.
    The pattern follows shell-style wildcard rules:
    * - matches any number of characters
    ? - matches a single character
    [seq] - matches any character in seq
    [!seq] - matches any character not in seq
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "pattern": (IO.STRING, {"default": "*.txt", "tooltip": "Shell-style pattern to match, e.g. '*.txt'."}),
            },
            "optional": {
                "recursive": (IO.BOOLEAN, {"default": False, "tooltip": "When True, '**' also matches inside subdirectories."}),
            },
        }

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("matching paths",)
    OUTPUT_TOOLTIPS = ("All paths matching the pattern, as a data list.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "glob_paths"
    OUTPUT_IS_LIST = (True,)

    # Class variable to store the last matched paths
    _last_matched_paths = {}

    @classmethod
    def IS_CHANGED(s, pattern: str, recursive: bool = False):
        # Get current paths
        current_paths = glob.glob(pattern, recursive=recursive)

        # Create a key for this specific pattern and recursive setting
        key = f"{pattern}_{recursive}"

        # If we haven't seen this pattern before, store it and trigger recalculation
        if key not in s._last_matched_paths:
            s._last_matched_paths[key] = current_paths
            return float("NaN")

        # Compare with previous paths
        previous_paths = s._last_matched_paths[key]
        if previous_paths != current_paths:
            # Update stored paths and trigger recalculation
            s._last_matched_paths[key] = current_paths
            return float("NaN")

        # No changes, return a consistent value
        import hashlib

        m = hashlib.md5()
        m.update(str(current_paths).encode())
        return m.hexdigest()

    def glob_paths(self, pattern: str, recursive: bool = False) -> tuple[list[str]]:
        return (glob.glob(pattern, recursive=recursive),)


class PathIsAbsolute(ComfyNodeABC):
    """
    Checks if a path is absolute.

    This node returns True if the path is absolute (begins at the root directory),
    and False if it's relative.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path to test."}),
            }
        }

    RETURN_TYPES = (IO.BOOLEAN,)
    RETURN_NAMES = ("is absolute",)
    OUTPUT_TOOLTIPS = ("True when the path is absolute (starts at the root).",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "check_is_absolute"

    def check_is_absolute(self, path: str) -> tuple[bool]:
        return (os.path.isabs(path),)


class PathIsDir(ComfyNodeABC):
    """
    Checks if a path points to a directory.

    This node returns True if the path exists and is a directory,
    and False otherwise.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path to test."}),
            }
        }

    RETURN_TYPES = (IO.BOOLEAN,)
    RETURN_NAMES = ("is dir",)
    OUTPUT_TOOLTIPS = ("True when the path exists and is a directory.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "check_is_dir"

    def check_is_dir(self, path: str) -> tuple[bool]:
        return (os.path.isdir(path),)


class PathIsFile(ComfyNodeABC):
    """
    Checks if a path points to a file.

    This node returns True if the path exists and is a regular file,
    and False otherwise.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path to test."}),
            }
        }

    RETURN_TYPES = (IO.BOOLEAN,)
    RETURN_NAMES = ("is file",)
    OUTPUT_TOOLTIPS = ("True when the path exists and is a regular file.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "check_is_file"

    def check_is_file(self, path: str) -> tuple[bool]:
        return (os.path.isfile(path),)


class PathJoin(ComfyNodeABC):
    """
    Joins multiple path components into a single path.

    This node takes multiple path components and joins them intelligently
    to form a single path. It handles directory separators correctly
    for the operating system.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path1": (IO.STRING, {"default": "", "tooltip": "First path component."}),
            },
            "optional": {
                "path2": (IO.STRING, {"default": "", "tooltip": "Second path component."}),
            },
        }

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("path",)
    OUTPUT_TOOLTIPS = ("The components joined into a single path.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "join_paths"

    def join_paths(self, path1: str, path2: str = "") -> tuple[str]:
        paths = [p for p in [path1, path2] if p]
        return (str(os.path.join(*paths)),)


class PathListDir(ComfyNodeABC):
    """
    Lists the contents of a directory.

    This node returns a list of files and directories in the specified path.
    If 'files_only' is True, it only returns files.
    If 'dirs_only' is True, it only returns directories.
    If both are False, it returns all contents.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Directory to list."}),
            },
            "optional": {
                "files_only": (IO.BOOLEAN, {"default": False, "tooltip": "When True, only files are returned."}),
                "dirs_only": (IO.BOOLEAN, {"default": False, "tooltip": "When True, only directories are returned."}),
            },
        }

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("entries",)
    OUTPUT_TOOLTIPS = ("The names of the directory entries, as a data list.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "list_directory"
    OUTPUT_IS_LIST = (True,)

    def list_directory(self, path: str, files_only: str = False, dirs_only: str = False) -> tuple[list[str]]:
        if not path:
            path = os.getcwd()

        if not os.path.exists(path):
            raise FileNotFoundError(f"Directory does not exist: {path}")
        if not os.path.isdir(path):
            raise NotADirectoryError(f"Basic data handling: Path is not a directory: {path}")

        entries = os.listdir(path)

        if files_only:
            entries = [e for e in entries if os.path.isfile(os.path.join(path, e))]
        elif dirs_only:
            entries = [e for e in entries if os.path.isdir(os.path.join(path, e))]

        return (entries,)


class PathNormalize(ComfyNodeABC):
    """
    Normalizes a path.

    This node normalizes a path by collapsing redundant separators,
    resolving up-level references, and converting to the correct
    separator for the operating system.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path to normalize."}),
            }
        }

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("normalized path",)
    OUTPUT_TOOLTIPS = ("The path with redundant separators and up-level references collapsed.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "normalize_path"

    def normalize_path(self, path: str) -> tuple[str]:
        return (os.path.normpath(path),)


class PathSetExtension(ComfyNodeABC):
    """
    Sets the file extension for a path.

    This node replaces the current extension in a path with a new one.
    The extension should include the dot (e.g., '.jpg').
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path whose extension is replaced."}),
                "extension": (IO.STRING, {"default": ".txt", "tooltip": "The new extension; a leading dot is added if missing."}),
            }
        }

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("path",)
    OUTPUT_TOOLTIPS = ("The path with its extension replaced.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "set_extension"

    def set_extension(self, path: str, extension: str) -> tuple[str]:
        # Make sure extension starts with a dot
        if not extension.startswith(".") and extension:
            extension = "." + extension

        root, _ = os.path.splitext(path)
        return (root + extension,)


class PathRelative(ComfyNodeABC):
    """
    Returns a relative path.

    This node computes a relative path from the 'start' path to the 'path'.
    If 'start' is not provided, the current working directory is used.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path to express relative to start."}),
            },
            "optional": {
                "start": (IO.STRING, {"default": "", "tooltip": "Base path; the current working directory is used when empty."}),
            },
        }

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("relative path",)
    OUTPUT_TOOLTIPS = ("The path expressed relative to start.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "get_relative_path"

    def get_relative_path(self, path: str, start: str = "") -> tuple[str]:
        if not start:
            start = os.getcwd()
        return (os.path.relpath(path, start),)


class PathSplit(ComfyNodeABC):
    """
    Splits a path into directory and filename components.

    This node takes a path and returns a tuple containing the directory path
    and the filename.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path to split into directory and filename."}),
            }
        }

    RETURN_TYPES = (IO.STRING, IO.STRING)
    RETURN_NAMES = ("directory", "filename")
    OUTPUT_TOOLTIPS = ("The directory (head) component.", "The filename (tail) component.")
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "split_path"

    def split_path(self, path: str) -> tuple[str, str]:
        return os.path.split(path)


class PathSplitExt(ComfyNodeABC):
    """
    Splits a path into name and extension components.

    This node takes a path and returns a tuple containing the path without
    the extension and the extension (including the dot).
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path to split into name and extension."}),
            }
        }

    RETURN_TYPES = (IO.STRING, IO.STRING)
    RETURN_NAMES = ("path without ext", "extension")
    OUTPUT_TOOLTIPS = ("The path without its extension.", "The extension including the dot (empty when none).")
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "split_ext"

    def split_ext(self, path: str) -> tuple[str, str]:
        return os.path.splitext(path)


class PathLoadStringFile(ComfyNodeABC):
    """
    Loads a text file in UTF-8 encoding and returns its content as a STRING
    without any further processing.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path of the UTF-8 text file to read."}),
            },
        }

    RETURN_TYPES = (IO.STRING, IO.BOOLEAN)
    RETURN_NAMES = ("text", "exists")
    OUTPUT_TOOLTIPS = ("The file content (empty when the file is missing or unreadable).", "True when the file exists and was read.")
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "load_text"

    @classmethod
    def IS_CHANGED(cls, path):
        try:
            if os.path.exists(path):
                return os.path.getmtime(path)
        except Exception:
            pass
        return float("NaN")  # Return NaN if file doesn't exist or can't access modification time

    def load_text(self, path: str):
        exists = os.path.exists(path)

        if not exists:
            return ("", False)

        try:
            with open(path, "r", encoding="utf-8") as f:
                text = f.read()
            return (text, True)
        except Exception:
            return ("", False)


class PathLoadImageRGB(ComfyNodeABC):
    """
    Loads an image from a file path and returns only the RGB channels.

    This node loads an image from the specified path and processes it to
    return only the RGB channels as a tensor, ignoring any alpha channel.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path of the image file to load."}),
            },
        }

    RETURN_TYPES = (IO.IMAGE, IO.BOOLEAN)
    RETURN_NAMES = ("image", "exists")
    OUTPUT_TOOLTIPS = ("The RGB image as a tensor (a blank 1x1 image when the file is missing).", "True when the image was loaded.")
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "load_image_rgb"

    @classmethod
    def IS_CHANGED(cls, path):
        try:
            if os.path.exists(path):
                return os.path.getmtime(path)
        except Exception:
            pass
        return float("NaN")  # Return NaN if file doesn't exist or can't access modification time

    def load_image_rgb(self, path: str):
        import numpy as np
        import torch

        img = load_image_helper(path)

        if img is None:
            # Create an empty 1x1 image
            empty_tensor = torch.zeros((1, 1, 1, 3), dtype=torch.float32)
            return (empty_tensor, False)

        # Convert to RGB (removing alpha if present)
        img_rgb = img.convert("RGB")

        # Convert to tensor format expected by ComfyUI
        image_tensor = np.array(img_rgb).astype(np.float32) / 255.0
        image_tensor = torch.from_numpy(image_tensor)[None,]

        return (image_tensor, True)


class PathLoadImageRGBA(ComfyNodeABC):
    """
    Loads an image from a file path and returns RGB channels and Alpha as a mask.

    This node loads an image from the specified path and processes it to
    return the RGB channels as a tensor and the Alpha channel as a mask tensor.
    If the image has no alpha channel, a blank mask is returned.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path of the image file to load."}),
            },
        }

    RETURN_TYPES = (IO.IMAGE, IO.MASK, IO.BOOLEAN)
    RETURN_NAMES = ("image", "mask", "exists")
    OUTPUT_TOOLTIPS = (
        "The RGB image as a tensor.",
        "The alpha channel as a mask (blank when the image has none).",
        "True when the image was loaded.",
    )
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "load_image_rgba"

    @classmethod
    def IS_CHANGED(cls, path):
        try:
            if os.path.exists(path):
                return os.path.getmtime(path)
        except Exception:
            pass
        return float("NaN")  # Return NaN if file doesn't exist or can't access modification time

    def load_image_rgba(self, path: str):
        import numpy as np
        import torch

        img = load_image_helper(path)

        if img is None:
            # Create empty 1x1 image and mask
            empty_image = torch.zeros((1, 1, 1, 3), dtype=torch.float32)
            empty_mask = torch.zeros((1, 1, 1), dtype=torch.float32)
            return (empty_image, empty_mask, False)

        # Convert to RGB for the image
        img_rgb = img.convert("RGB")

        # Convert to tensor format expected by ComfyUI
        image_tensor = np.array(img_rgb).astype(np.float32) / 255.0
        image_tensor = torch.from_numpy(image_tensor)[None,]

        # Extract alpha channel as mask
        mask_tensor = extract_mask_from_alpha(img)

        return (image_tensor, mask_tensor, True)


class PathLoadMaskFromAlpha(ComfyNodeABC):
    """
    Loads a mask from the alpha channel of an image.

    This node loads an image from the specified path and extracts the alpha
    channel to use as a mask. If the image has no alpha channel, a blank mask
    is returned.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path of the image whose alpha channel is used."}),
            },
        }

    RETURN_TYPES = (IO.MASK, IO.BOOLEAN)
    RETURN_NAMES = ("mask", "exists")
    OUTPUT_TOOLTIPS = ("The alpha channel as a mask (blank when the image has none or is missing).", "True when the image was loaded.")
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "load_mask_from_alpha"

    @classmethod
    def IS_CHANGED(cls, path):
        try:
            if os.path.exists(path):
                return os.path.getmtime(path)
        except Exception:
            pass
        return float("NaN")  # Return NaN if file doesn't exist or can't access modification time

    def load_mask_from_alpha(self, path: str):
        import torch

        img = load_image_helper(path)

        if img is None:
            # Return empty 1x1 mask
            empty_mask = torch.zeros((1, 1, 1), dtype=torch.float32)
            return (empty_mask, False)

        mask_tensor = extract_mask_from_alpha(img)
        return (mask_tensor, True)


class PathLoadMaskFromGreyscale(ComfyNodeABC):
    """
    Loads a mask from a greyscale image or the red channel of an RGB image.

    This node loads an image from the specified path and creates a mask from it.
    If the image is greyscale, the intensity is used directly.
    If the image is RGB, the red channel is used.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": (IO.STRING, {"default": "", "tooltip": "Path of the image to build the mask from."}),
            },
            "optional": {
                "invert": (IO.BOOLEAN, {"default": False, "tooltip": "Invert the mask (1.0 - mask) after extraction."}),
            },
        }

    RETURN_TYPES = (IO.MASK, IO.BOOLEAN)
    RETURN_NAMES = ("mask", "exists")
    OUTPUT_TOOLTIPS = ("The mask derived from the greyscale/red channel.", "True when the image was loaded.")
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "load_mask_from_greyscale"

    @classmethod
    def IS_CHANGED(cls, path):
        try:
            if os.path.exists(path):
                return os.path.getmtime(path)
        except Exception:
            pass
        return float("NaN")  # Return NaN if file doesn't exist or can't access modification time

    def load_mask_from_greyscale(self, path: str, invert: bool = False):
        import torch

        img = load_image_helper(path)

        if img is None:
            # Return empty 1x1 mask
            empty_mask = torch.zeros((1, 1, 1), dtype=torch.float32)
            return (empty_mask, False)

        mask_tensor = extract_mask_from_greyscale(img)

        # Optionally invert the mask (1.0 - mask)
        if invert:
            mask_tensor = 1.0 - mask_tensor

        return (mask_tensor, True)


class PathSaveStringFile(ComfyNodeABC):
    """
    Saves a string to a text file.

    This node takes a string and saves it to the specified path as a text file.
    Optionally, you can choose to create the directory if it doesn't exist.
    If 'append' is True, the text is appended to an existing file instead of
    overwriting it.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text": (IO.STRING, {"default": "", "tooltip": "The text to write."}),
                "path": (IO.STRING, {"default": "", "tooltip": "Destination file path."}),
            },
            "optional": {
                "create_dirs": (IO.BOOLEAN, {"default": True, "tooltip": "Create missing parent directories."}),
                "append": (IO.BOOLEAN, {"default": False, "tooltip": "Append to an existing file instead of overwriting it."}),
                "encoding": (IO.STRING, {"default": "utf-8", "tooltip": "Text encoding to use when writing."}),
            },
        }

    RETURN_TYPES = (IO.BOOLEAN,)
    RETURN_NAMES = ("success",)
    OUTPUT_TOOLTIPS = ("True when the file was written successfully.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "save_text"
    OUTPUT_NODE = True

    def save_text(self, text: str, path: str, create_dirs: bool = True, append: bool = False, encoding: str = "utf-8"):
        if not path:
            print("Basic data handling: Save failed - no path specified")
            return (False,)

        try:
            # Create directories if needed
            directory = os.path.dirname(path)
            if directory and create_dirs and not os.path.exists(directory):
                os.makedirs(directory)

            mode = "a" if append else "w"
            with open(path, mode, encoding=encoding) as f:
                f.write(text)

            action = "appended" if append else "saved"
            print(f"Basic data handling: Successfully {action} text to {path}")
            return (True,)
        except Exception as e:
            print(f"Basic data handling: Error saving text file: {e}")
            return (False,)


def compose_prompt_text(prompt: str, negative_prompt: str) -> str:
    """
    Build the generation-parameter text embedded into saved images.

    Follows the Stable Diffusion WebUI convention: the (optional) positive
    prompt is written first, followed by an optional ``Negative prompt:``
    line::

        <positive prompt>
        Negative prompt: <negative prompt>

    Returns an empty string when neither value is provided, in which case no
    metadata is embedded into the file.
    """
    lines = []
    if prompt.strip():
        lines.append(prompt.strip())
    if negative_prompt.strip():
        lines.append(f"Negative prompt: {negative_prompt.strip()}")
    return "\n".join(lines)


def build_png_info(metadata_text: str):
    """
    Wrap ``metadata_text`` in a Pillow ``PngInfo`` container under the standard
    ``parameters`` text-chunk key so it can be embedded in a PNG file.

    Returns ``None`` when there is no text to embed (or Pillow's PNG metadata
    support is unavailable), in which case the image should be saved without
    extra metadata.
    """
    if not metadata_text:
        return None
    try:
        from PIL import PngImagePlugin
    except ModuleNotFoundError:
        return None
    pnginfo = PngImagePlugin.PngInfo()
    pnginfo.add_text("parameters", metadata_text)
    return pnginfo


def build_image_exif(metadata_text: str, include_description: bool = True):
    """
    Build an EXIF block that stores ``metadata_text`` for formats without a
    native text chunk (JPEG, WEBP, JXL).

    The payload is written into the EXIF ``UserComment`` field (tag 0x9286) of
    the Exif IFD as ``UNICODE\0`` + UTF-16-BE, which matches what Stable
    Diffusion WebUI / piexif based readers expect. When ``include_description``
    is true, the payload is also written as UTF-8 into the EXIF
    ``ImageDescription`` field (tag 0x010E) of IFD0.

    Returns the EXIF bytes (starting with the ``Exif\0\0`` marker), or ``None``
    when there is no text to embed.
    """
    if not metadata_text:
        return None
    try:
        from PIL import ExifTags
    except ModuleNotFoundError:
        return None
    Image, _ = _require_pillow()
    exif = Image.Exif()
    if include_description:
        exif[0x010E] = metadata_text.encode("utf-8")
    exif.get_ifd(ExifTags.IFD.Exif)[0x9286] = b"UNICODE\x00" + metadata_text.encode("utf-16-be")
    return exif.tobytes()


def build_xmp_packet(metadata_text: str) -> bytes:
    """
    Build an XMP packet storing ``metadata_text`` in the Dublin Core
    ``dc:description`` tag, as expected for JPEG XL ``xml `` boxes.
    """
    from xml.sax.saxutils import escape

    body = escape(metadata_text)
    packet = (
        '<?xpacket begin="\ufeff" id="W5M0MpCehiHzreSzNTczkc9d"?>\n'
        '<x:xmpmeta xmlns:x="adobe:ns:meta/">\n'
        '<rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">\n'
        '<rdf:Description xmlns:dc="http://purl.org/dc/elements/1.1/">\n'
        '<dc:description><rdf:Alt><rdf:li xml:lang="x-default">' + body + "</rdf:li></rdf:Alt></dc:description>\n"
        "</rdf:Description>\n"
        "</rdf:RDF>\n"
        "</x:xmpmeta>\n"
        '<?xpacket end="w"?>'
    )
    return packet.encode("utf-8")


def metadata_save_kwargs(metadata_text: str, fmt: str) -> dict:
    """
    Return the extra keyword arguments that embed ``metadata_text`` when saving
    an image in the (lower-case) format ``fmt``.

    Returns an empty dict when there is no text to embed or when the format
    cannot carry text metadata.
    """
    if not metadata_text:
        return {}
    if fmt == "png":
        return {"pnginfo": build_png_info(metadata_text)}
    if fmt in ("jpg", "jpeg"):
        exif = build_image_exif(metadata_text, include_description=True)
        return {"exif": exif} if exif is not None else {}
    if fmt in ("webp", "jxl"):
        exif = build_image_exif(metadata_text, include_description=False)
        kwargs = {"exif": exif} if exif is not None else {}
        if fmt == "jxl":
            # EXIF and XMP boxes are only available in the JXL container format
            kwargs["use_container"] = True
            kwargs["xmp"] = build_xmp_packet(metadata_text)
        return kwargs
    return {}


# --- shared helpers for the image-save nodes --------------------------------
#
# The save-format dropdown is not hard-coded: it is derived from the installed
# Pillow (+ any image plugins such as pillow-jxl it can save with), so it always
# matches what the running library supports and automatically picks up new
# formats/plugins. Only "formats" that are not meaningful photographic raster
# output (icons / GPU textures / vector / multi-image containers, e.g. DDS,
# ICNS, EPS, MPO) are never offered even when Pillow can technically write them.
#
# A save *token* is the short, user-facing format name and equals the file
# extension we write (jpg/tif/jp2/...). Each token maps to the Pillow format-id
# used with ``Image.save(format=...)``.
_PIL_FORMAT_BY_TOKEN: dict[str, str] = {
    "png": "PNG",
    "jpg": "JPEG",
    "jpeg": "JPEG",  # alias; normalised to "jpg"
    "webp": "WEBP",
    "jxl": "JXL",
    "bmp": "BMP",
    "tif": "TIFF",
    "tiff": "TIFF",  # alias; normalised to "tif"
    "gif": "GIF",
    "jp2": "JPEG2000",
    "jpeg2000": "JPEG2000",  # alias; normalised to "jp2"
    "ppm": "PPM",
    "pcx": "PCX",
    "tga": "TGA",
    "qoi": "QOI",
    "avif": "AVIF",
}

# friendly ordering for the dropdown: the well-known set first, then whatever
# else is discovered in the installed library/plugins.
_IMAGE_SAVE_ORDER = ["png", "jpg", "webp", "jxl", "bmp", "tif", "gif", "jp2", "ppm", "pcx", "tga", "qoi", "avif"]

# Pillow format-ids that are not meaningful photographic raster *stills* for this
# save node (icons/cursors, GPU textures, vector/postscript, single-multi-image
# containers, scientific dumps). Never offered even if Pillow can write them.
_NON_PHOTOGRAPHIC_RASTER_IDS = {
    "ICO",
    "CUR",
    "ICNS",
    "DDS",
    "DIB",
    "EPS",
    "IM",
    "MPO",
    "SGI",
    "PDF",
    "WMF",
    "EMF",
    "BUFR",
    "GRIB",
    "HDF5",
    "FITS",
}

# canonical token for a couple of common user-spelled aliases
_FORMAT_ALIASES = {"jpeg": "jpg", "tiff": "tif", "jpeg2000": "jp2"}


def _load_image_plugins() -> None:
    """Import optional Pillow plugins (e.g. pillow-jxl) so Pillow knows about
    every format the library can actually save."""
    try:
        import pillow_jxl  # noqa: F401 - imported but unused, registers the JXL plugin
    except ModuleNotFoundError:
        pass


def _writable_pillow_formats(mode: str) -> set[str]:
    """
    The set of Pillow format-ids that can successfully encode a small image in
    ``mode`` (probed at runtime against the actually installed library/plugins).
    """
    import io as _io
    from PIL import Image as _PIL

    _load_image_plugins()
    ids = set(_PIL.SAVE.keys())
    ids.update(str(v).upper() for v in _PIL.registered_extensions().values())
    writable: set[str] = set()
    for fid in sorted(ids):
        try:
            buf = _io.BytesIO()
            _PIL.new(mode, (4, 4)).save(buf, format=fid)
            buf.seek(0)
            _PIL.open(buf).load()
            writable.add(str(fid).upper())
        except Exception:
            continue
    return writable


def _discover_save_formats() -> tuple[list[str], list[str]]:
    """
    Derive ``(all_formats, alpha_formats)`` dropdown lists from Pillow/plugins.
    ``alpha_formats`` (surfaced on the IMAGE+MASK node) is the subset that can
    also hold an alpha channel.
    """
    writable_rgb = _writable_pillow_formats("RGB")
    writable_rgba = _writable_pillow_formats("RGBA")

    # 1) tokens we already know about and that are actually writable
    offered = [t for t in _IMAGE_SAVE_ORDER if _PIL_FORMAT_BY_TOKEN[t] in writable_rgb]
    known_ids = {_PIL_FORMAT_BY_TOKEN[t] for t in offered}

    # 2) any additional writable photographic format in the library/plugins
    extras = []
    for pid in sorted(writable_rgb):
        if pid in known_ids or pid in _NON_PHOTOGRAPHIC_RASTER_IDS:
            continue
        token = pid.lower()
        # register it so it can be written through the same machinery
        _PIL_FORMAT_BY_TOKEN[token] = pid
        if token not in _IMAGE_SAVE_ORDER:
            _IMAGE_SAVE_ORDER.append(token)
        extras.append(token)

    all_formats = offered + extras
    alpha_formats = [t for t in all_formats if _PIL_FORMAT_BY_TOKEN[t] in writable_rgba]
    return all_formats, alpha_formats


# Discovering the formats probes the installed Pillow/plugins, so it must not be
# required for importing this module: environments that only inspect the node
# metadata (e.g. the comfy-org/node-diff CI) have no Pillow installed. Fall back
# to the classic formats there instead of failing to load the whole node pack.
try:
    (_IMAGE_SAVE_FORMATS, _IMAGE_SAVE_ALPHA_FORMATS) = _discover_save_formats()
except ModuleNotFoundError:
    _IMAGE_SAVE_FORMATS = ["png", "jpg", "webp", "jxl", "bmp", "tif", "gif"]
    _IMAGE_SAVE_ALPHA_FORMATS = ["png", "webp", "jxl"]
else:
    # Always make sure the most common formats are present when the backend
    # supports them (guards against an odd Pillow not probing its core formats).
    for _token in ("png", "jpg", "webp"):
        if _token in _PIL_FORMAT_BY_TOKEN and _PIL_FORMAT_BY_TOKEN[_token] and _token not in _IMAGE_SAVE_FORMATS:
            _IMAGE_SAVE_FORMATS.insert(0, _token)

    # Log the supported save formats once at node-load time, so users can see
    # what the running Pillow/image plugins can write.
    print("Basic data handling: supported image formats: " + ", ".join(_IMAGE_SAVE_FORMATS))
    print("Basic data handling: supported IMAGE+MASK (alpha) formats: " + ", ".join(_IMAGE_SAVE_ALPHA_FORMATS))


def _normalize_image_format(format: str) -> str:
    """
    Normalise a user-supplied format to its canonical save token (lower-cased),
    mapping common aliases (jpeg->jpg, tiff->tif, jpeg2000->jp2) so both the
    dropdown spellings and values wired in from elsewhere are accepted. Unknown
    tokens are passed through unchanged (they may name a format the writer will
    then attempt, or reject gracefully).
    """
    token = str(format or "png").strip().lower()
    return _FORMAT_ALIASES.get(token, token)


def _comfy_format_date(fmt_text: str, when) -> str:
    """
    Format ``when`` using ComfyUI's own ``formatDate`` semantics (used by
    ``%date:<format>%``). Supported tokens: ``yyyy``/``yy`` for the year, and
    ``M``/``MM``, ``d``/``dd``, ``h``/``hh``, ``m``/``mm``, ``s``/``ss``
    (lower-case ``h`` for hours, exactly like ComfyUI). Single letters are
    unpadded, doubled letters are zero-padded.
    """
    import re

    parts = {
        "d": when.day,
        "M": when.month,
        "h": when.hour,
        "m": when.minute,
        "s": when.second,
    }
    # Mirrors ComfyUI's regex: dd?|MM?|hh?|mm?|ss?|yyy?y? (yyy?y? last so a
    # four-digit year is matched as a single token, not two "yy" groups).
    pattern = re.compile(r"dd?|MM?|hh?|mm?|ss?|yyy?y?")

    def replace(token: str) -> str:
        if token == "yy":
            return str(when.year)[-2:]
        if token == "yyyy":
            return str(when.year)
        if token and token[0] in parts:
            # padStart(token.length, "0")
            return str(parts[token[0]]).zfill(len(token))
        return token

    return pattern.sub(lambda m: replace(m.group(0)), fmt_text)


def _expand_filename_tokens(text: str, width: int = 0, height: int = 0) -> str:
    """
    Expand ComfyUI-style filename templates in ``text`` for both the plain path
    and the ``filename_prefix`` modes.

    Handles ``%date:<format>%`` (using ComfyUI's own date formatting) as well as
    the standard built-in tokens ``%width%``, ``%height%``, ``%year%``,
    ``%month%``, ``%day%``, ``%hour%``, ``%minute%`` and ``%second%``. Node
    reference tokens (``%Node.widget%``) can only be resolved by the ComfyUI
    frontend, so they are left untouched here.
    """
    import re
    from datetime import datetime

    if not text:
        return text

    now = datetime.now()
    text = re.sub(r"%date:(.*?)%", lambda m: _comfy_format_date(m.group(1), now), text)
    text = text.replace("%width%", str(width))
    text = text.replace("%height%", str(height))
    text = text.replace("%year%", str(now.year))
    text = text.replace("%month%", str(now.month).zfill(2))
    text = text.replace("%day%", str(now.day).zfill(2))
    text = text.replace("%hour%", str(now.hour).zfill(2))
    text = text.replace("%minute%", str(now.minute).zfill(2))
    text = text.replace("%second%", str(now.second).zfill(2))
    return text


def _official_save_image_path(filename_prefix, output_dir, image_width=0, image_height=0):
    """
    Resolve ``filename_prefix`` to concrete output folder/filename/counter using
    ComfyUI's own ``folder_paths.get_save_image_path`` wherever possible.

    Returns ``(full_output_folder, filename, counter, subfolder, resolved_prefix)``
    exactly like the official function. When ``folder_paths`` is unavailable
    (standalone test-suite / non-ComfyUI host), a minimal local equivalent that
    mirrors the same subfolder + ``_<counter>_`` scheme is used instead.
    """
    if _official_get_save_image_path is not None:
        return _official_get_save_image_path(filename_prefix, output_dir, image_width, image_height)

    # --- fallback (only reached outside ComfyUI) -------------------------
    def expand(text, w, h):
        import time

        now = time.localtime()
        text = text.replace("%width%", str(w)).replace("%height%", str(h))
        text = text.replace("%year%", str(now.tm_year))
        text = text.replace("%month%", str(now.tm_mon).zfill(2))
        text = text.replace("%day%", str(now.tm_mday).zfill(2))
        text = text.replace("%hour%", str(now.tm_hour).zfill(2))
        text = text.replace("%minute%", str(now.tm_min).zfill(2))
        text = text.replace("%second%", str(now.tm_sec).zfill(2))
        return text

    if "%" in filename_prefix:
        filename_prefix = expand(filename_prefix, image_width, image_height)
    subfolder = os.path.dirname(os.path.normpath(filename_prefix))
    filename = os.path.basename(os.path.normpath(filename_prefix))
    full_output_folder = os.path.join(output_dir, subfolder)

    highest = 0
    try:
        for entry in os.listdir(full_output_folder):
            stem, _, _ = entry.rpartition(".")
            if stem.startswith(filename + "_") and stem[len(filename) + 1 :].rstrip("_").isdigit():
                digits = stem[len(filename) + 1 :].rstrip("_")
                if digits.isdigit():
                    highest = max(highest, int(digits))
    except FileNotFoundError:
        os.makedirs(full_output_folder, exist_ok=True)
    return full_output_folder, filename, highest + 1, subfolder, filename_prefix


def _plan_save_paths(path: str, format: str, use_prefix_mode: bool, frame_count: int, width: int, height: int) -> tuple[list[str], bool]:
    """
    Compute the absolute destination path(s) for ``frame_count`` images.

    Returns ``(paths, create_dirs)``:
    - ``use_prefix_mode`` False: ``path`` is a concrete file location. A single
      frame writes exactly to ``path`` (+ extension); several frames get an
      incrementing ``_00000``-style suffix so they do not overwrite each other.
    - ``use_prefix_mode`` True: ``path`` is a ComfyUI ``filename_prefix`` under
      the output folder, resolved through ComfyUI's own filename helpers and
      auto-numbered like the built-in "Save Image" node.
    """
    fmt = _normalize_image_format(format)
    # Expand ComfyUI-style templates (%date:...%, %width%, %height%, ...) up
    # front so they apply in both the plain-path and the prefix modes.
    path = _expand_filename_tokens(path, width, height)

    if use_prefix_mode:
        full_output_folder, filename, counter, _, _ = _official_save_image_path(path, get_output_directory(), width, height)
        paths = []
        for _ in range(frame_count):
            name = f"{filename}_{counter:05}_.{fmt}"
            paths.append(os.path.join(full_output_folder, name))
            counter += 1
        return paths, False

    if not path.endswith(f".{fmt}"):
        base = f"{path}.{fmt}"
    else:
        base = path

    if frame_count <= 1:
        return [base], True

    root, ext = os.path.splitext(base)
    paths = [f"{root}_{i:05}{ext}" for i in range(frame_count)]
    return paths, True


def _ensure_parent_directories(path: str, create_dirs: bool) -> None:
    """Create the parent directory(ies) of ``path`` when requested."""
    directory = os.path.dirname(path)
    if directory and create_dirs and not os.path.exists(directory):
        os.makedirs(directory)


def _has_jxl_support() -> bool:
    """Return True when the pillow-jxl plugin is importable."""
    try:
        import pillow_jxl  # noqa: F401 - imported but unused, kept for JPEG XL support

        return True
    except ModuleNotFoundError:
        return False


def _write_pil_image(pil_img, path: str, fmt: str, quality: int, metadata_text: str) -> bool:
    """
    Write ``pil_img`` to ``path``. ``fmt`` is a canonical save token (see
    ``_PIL_FORMAT_BY_TOKEN``). For the lossy/metadata-capable set (png/jpg/webp/
    jxl) the prompt text is embedded as ``parameters`` metadata and ``quality``
    is honoured; all other discovered formats are written plainly. Returns True
    on success.
    """
    try:
        if fmt == "png":
            pil_img.save(path, format="PNG", **metadata_save_kwargs(metadata_text, fmt))
        elif fmt == "jpg":
            # JPEG cannot carry alpha; drop it so an RGBA input still saves.
            pil_img.convert("RGB").save(path, format="JPEG", quality=quality, **metadata_save_kwargs(metadata_text, fmt))
        elif fmt == "webp":
            pil_img.save(path, format="WEBP", quality=quality, **metadata_save_kwargs(metadata_text, fmt))
        elif fmt == "jxl":
            pil_img.save(path, format="JXL", quality=quality, **metadata_save_kwargs(metadata_text, fmt))
        else:
            if metadata_text:
                print("Basic data handling: Prompt metadata is not supported for this format; skipping it.")
            pil_img.save(path, format=_PIL_FORMAT_BY_TOKEN.get(fmt, fmt.upper()))
        return True
    except Exception as e:
        print(f"Basic data handling: Error saving image: {e}")
        return False


class PathSaveImageRGB(ComfyNodeABC):
    """
    Saves an image to a file.

    This node takes an image tensor and saves it to the specified path.
    Supports various image formats like PNG, JPG, WEBP, JXL (if pillow-jxl is installed), etc.

    By default ``path`` is a concrete absolute/relative file location and every
    image frame of the batch is written (several frames get an incrementing
    suffix). When ``use_prefix_mode`` is enabled ``path`` is instead treated as
    a ComfyUI ``filename_prefix`` under the output folder, named and
    auto-numbered exactly like the built-in "Save Image" node.

    When ``prompt`` and/or ``negative_prompt`` are provided, they are embedded
    into the saved image as ``parameters`` metadata: in the PNG text chunk, in
    the EXIF ``UserComment`` (and ``ImageDescription`` for JPEG) fields, and in
    the EXIF + XMP boxes for JPEG XL. Formats that cannot carry text metadata
    ignore the prompts.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": (IO.IMAGE,),
                "path": (
                    IO.STRING,
                    {
                        "default": "",
                        "tooltip": 'Destination file path (an extension is added from the format when missing), or a ComfyUI filename_prefix under the output folder when "use prefix mode" is enabled.',
                    },
                ),
            },
            "optional": {
                "format": (
                    _IMAGE_SAVE_FORMATS,
                    {
                        "default": "png",
                        "tooltip": "Image save format (png, jpg, webp, bmp, ...). The list is derived from what the installed Pillow/plugins can write. Drag a STRING onto this to override it.",
                    },
                ),
                "quality": (IO.INT, {"default": 95, "min": 1, "max": 100, "tooltip": "Quality for lossy formats (jpg/webp/jxl)."}),
                "create_dirs": (IO.BOOLEAN, {"default": True, "tooltip": "Create missing parent directories (plain path mode only)."}),
                "prompt": (IO.STRING, {"default": "", "tooltip": "Optional positive prompt embedded as parameters metadata."}),
                "negative_prompt": (IO.STRING, {"default": "", "tooltip": "Optional negative prompt embedded as parameters metadata."}),
                # Appended last on purpose: the node stores widget values
                # positionally, so adding before the existing widgets would shift
                # old workflows and mis-assign their saved values.
                "use_prefix_mode": (
                    IO.BOOLEAN,
                    {
                        "default": False,
                        "tooltip": "When True, 'path' is treated as a ComfyUI filename_prefix under the output folder and files are named/auto-numbered exactly like ComfyUI's \"Save Image\" node.",
                    },
                ),
            },
        }

    RETURN_TYPES = (IO.BOOLEAN,)
    RETURN_NAMES = ("success",)
    OUTPUT_TOOLTIPS = ("True when the image was saved successfully.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "save_image"
    OUTPUT_NODE = True

    def save_image(
        self,
        images,
        path: str = "",
        use_prefix_mode: bool = False,
        format: str = "png",
        quality: int = 95,
        create_dirs: bool = True,
        prompt: str = "",
        negative_prompt: str = "",
    ):
        if not path:
            print("Basic data handling: Save failed - no path specified")
            return (False,)

        import numpy as np
        from PIL import Image

        fmt = _normalize_image_format(format)
        if fmt == "jxl" and not _has_jxl_support():
            print(
                "Basic data handling: JPEG XL format requested but pillow_jxl module is not installed. "
                "Please install it with 'pip install pillow-jxl-plugin'."
            )
            return (False,)

        batch = len(images)
        height, width = images.shape[1], images.shape[2]
        paths, create_dirs = _plan_save_paths(path, fmt, use_prefix_mode, batch, width, height)

        # Compose the prompt metadata to embed into the saved file once.
        metadata_text = compose_prompt_text(prompt, negative_prompt)

        for index, target_path in enumerate(paths):
            _ensure_parent_directories(target_path, create_dirs)

            # Convert from tensor format back to PIL Image
            img_tensor = images[index].cpu().numpy()
            img_np = (img_tensor * 255).astype(np.uint8)
            pil_img = Image.fromarray(img_np)

            if not _write_pil_image(pil_img, target_path, fmt, quality, metadata_text):
                return (False,)
            print(f"Basic data handling: Successfully saved image to {target_path}")

        return (True,)


class PathSaveImageRGBA(ComfyNodeABC):
    """
    Saves an image with a mask to a file with transparency.

    This node takes an image tensor and a mask tensor and saves them to the
    specified path as an image with transparency, where the mask defines the
    alpha channel.

    By default ``path`` is a concrete absolute/relative file location and every
    image frame of the batch is written (several frames get an incrementing
    suffix). When ``use_prefix_mode`` is enabled ``path`` is instead treated as
    a ComfyUI ``filename_prefix`` under the output folder, named and
    auto-numbered exactly like the built-in "Save Image" node.

    When ``prompt`` and/or ``negative_prompt`` are provided, they are embedded
    into the saved image as ``parameters`` metadata: in the PNG text chunk, in
    the EXIF ``UserComment`` (and ``ImageDescription`` for JPEG) fields, and in
    the EXIF + XMP boxes for JPEG XL. Formats that cannot carry text metadata
    ignore the prompts.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": (IO.IMAGE,),
                "mask": (IO.MASK,),
                "path": (
                    IO.STRING,
                    {
                        "default": "",
                        "tooltip": 'Destination file path (an extension is added from the format when missing), or a ComfyUI filename_prefix under the output folder when "use prefix mode" is enabled.',
                    },
                ),
            },
            "optional": {
                "format": (
                    _IMAGE_SAVE_ALPHA_FORMATS,
                    {
                        "default": "png",
                        "tooltip": "Image save format that can hold an alpha channel (png, webp, ...; jpg is coerced to png). Derived from what the installed Pillow/plugins can write. Drag a STRING onto this to override it.",
                    },
                ),
                "quality": (IO.INT, {"default": 95, "min": 1, "max": 100, "tooltip": "Quality for lossy formats (webp/jxl)."}),
                "invert_mask": (IO.BOOLEAN, {"default": False, "tooltip": "Invert the mask before using it as the alpha channel."}),
                "create_dirs": (IO.BOOLEAN, {"default": True, "tooltip": "Create missing parent directories (plain path mode only)."}),
                "prompt": (IO.STRING, {"default": "", "tooltip": "Optional positive prompt embedded as parameters metadata."}),
                "negative_prompt": (IO.STRING, {"default": "", "tooltip": "Optional negative prompt embedded as parameters metadata."}),
                # Appended last on purpose: the node stores widget values
                # positionally, so adding before the existing widgets would shift
                # old workflows and mis-assign their saved values.
                "use_prefix_mode": (
                    IO.BOOLEAN,
                    {
                        "default": False,
                        "tooltip": "When True, 'path' is treated as a ComfyUI filename_prefix under the output folder and files are named/auto-numbered exactly like ComfyUI's \"Save Image\" node.",
                    },
                ),
            },
        }

    RETURN_TYPES = (IO.BOOLEAN,)
    RETURN_NAMES = ("success",)
    OUTPUT_TOOLTIPS = ("True when the image with alpha was saved successfully.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "save_image_with_mask"
    OUTPUT_NODE = True

    def save_image_with_mask(
        self,
        images,
        mask,
        path: str = "",
        use_prefix_mode: bool = False,
        format: str = "png",
        quality: int = 95,
        invert_mask: bool = False,
        create_dirs: bool = True,
        prompt: str = "",
        negative_prompt: str = "",
    ):
        if not path:
            print("Basic data handling: Save failed - no path specified")
            return (False,)

        import numpy as np
        from PIL import Image

        # JPEG doesn't support transparency -> coerce to PNG, as before.
        fmt = _normalize_image_format(format)
        if fmt in ("jpg", "jpeg"):
            print("Basic data handling: JPEG format doesn't support transparency. Using PNG instead.")
            fmt = "png"
        if fmt == "jxl" and not _has_jxl_support():
            print(
                "Basic data handling: JPEG XL format requested but pillow_jxl module is not installed. "
                "Please install it with 'pip install pillow-jxl-plugin'."
            )
            return (False,)

        batch = len(images)
        height, width = images.shape[1], images.shape[2]
        paths, create_dirs = _plan_save_paths(path, fmt, use_prefix_mode, batch, width, height)

        mask_batch = len(mask)
        # Mask tensor may be padded to the image batches; round-robin when the
        # mask has fewer frames (e.g. a single mask applied to every frame).
        mask_frames = mask.cpu()

        metadata_text = compose_prompt_text(prompt, negative_prompt)

        for index, target_path in enumerate(paths):
            _ensure_parent_directories(target_path, create_dirs)

            img_tensor = images[index].cpu().numpy()
            mask_tensor = mask_frames[index % mask_batch]

            # Invert the mask if needed (1.0 becomes transparent, 0.0 becomes opaque)
            if invert_mask:
                mask_tensor = 1.0 - mask_tensor

            # Convert to alpha channel (0-255)
            alpha_np = (255.0 * (1.0 - mask_tensor.numpy())).astype(np.uint8)

            # Convert to uint8 format for PIL
            img_np = (img_tensor * 255).astype(np.uint8)

            # Create PIL image (RGB)
            pil_img = Image.fromarray(img_np)

            # Create alpha channel image (avoid deprecated 'mode' kwarg in Pillow 13+)
            alpha_img = Image.fromarray(alpha_np).convert("L")

            # Convert to RGBA and add alpha channel
            pil_img_rgba = pil_img.convert("RGBA")
            pil_img_rgba.putalpha(alpha_img)

            if not _write_pil_image(pil_img_rgba, target_path, fmt, quality, metadata_text):
                return (False,)
            print(f"Basic data handling: Successfully saved image with mask to {target_path}")

        return (True,)


class PathInputDir(ComfyNodeABC):
    """
    Returns the ComfyUI input path.

    This is where input images are usually stored when using ComfyUI
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {}

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("input_path",)
    OUTPUT_TOOLTIPS = ("Absolute path of ComfyUI's input directory.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "execute"
    DISPLAY_NAME = "Get ComfyUI Input Path"

    def execute(self) -> tuple[str]:
        return (get_input_directory(),)


class PathOutputDir(ComfyNodeABC):
    """
    Returns the ComfyUI output path.

    This is where output images are usually stored when using ComfyUI
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {}

    RETURN_TYPES = (IO.STRING,)
    RETURN_NAMES = ("output_path",)
    OUTPUT_TOOLTIPS = ("Absolute path of ComfyUI's output directory.",)
    CATEGORY = "Basic/Path"
    DESCRIPTION = cleandoc(__doc__ or "")
    FUNCTION = "execute"
    DISPLAY_NAME = "Get ComfyUI Output Path"

    def execute(self) -> tuple[str]:
        return (get_output_directory(),)


NODE_CLASS_MAPPINGS = {
    "Basic data handling: PathAbspath": PathAbspath,
    "Basic data handling: PathBasename": PathBasename,
    "Basic data handling: PathCommonPrefix": PathCommonPrefix,
    "Basic data handling: PathDirname": PathDirname,
    "Basic data handling: PathExists": PathExists,
    "Basic data handling: PathExpandVars": PathExpandVars,
    "Basic data handling: PathGetCwd": PathGetCwd,
    "Basic data handling: PathGetExtension": PathGetExtension,
    "Basic data handling: PathSetExtension": PathSetExtension,
    "Basic data handling: PathGetSize": PathGetSize,
    "Basic data handling: PathGlob": PathGlob,
    "Basic data handling: PathInputDir": PathInputDir,
    "Basic data handling: PathIsAbsolute": PathIsAbsolute,
    "Basic data handling: PathIsDir": PathIsDir,
    "Basic data handling: PathIsFile": PathIsFile,
    "Basic data handling: PathJoin": PathJoin,
    "Basic data handling: PathListDir": PathListDir,
    "Basic data handling: PathNormalize": PathNormalize,
    "Basic data handling: PathOutputDir": PathOutputDir,
    "Basic data handling: PathRelative": PathRelative,
    "Basic data handling: PathSplit": PathSplit,
    "Basic data handling: PathSplitExt": PathSplitExt,
    "Basic data handling: PathLoadStringFile": PathLoadStringFile,
    "Basic data handling: PathLoadImageRGB": PathLoadImageRGB,
    "Basic data handling: PathLoadImageRGBA": PathLoadImageRGBA,
    "Basic data handling: PathLoadMaskFromAlpha": PathLoadMaskFromAlpha,
    "Basic data handling: PathLoadMaskFromGreyscale": PathLoadMaskFromGreyscale,
    "Basic data handling: PathSaveStringFile": PathSaveStringFile,
    "Basic data handling: PathSaveImageRGB": PathSaveImageRGB,
    "Basic data handling: PathSaveImageRGBA": PathSaveImageRGBA,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "Basic data handling: PathAbspath": "abspath",
    "Basic data handling: PathBasename": "basename",
    "Basic data handling: PathCommonPrefix": "common prefix",
    "Basic data handling: PathDirname": "dirname",
    "Basic data handling: PathExists": "exists",
    "Basic data handling: PathExpandVars": "expand vars",
    "Basic data handling: PathGetCwd": "get current working directory",
    "Basic data handling: PathGetExtension": "get extension",
    "Basic data handling: PathSetExtension": "set extension",
    "Basic data handling: PathGetSize": "get size",
    "Basic data handling: PathGlob": "glob",
    "Basic data handling: PathInputDir": "input dir",
    "Basic data handling: PathIsAbsolute": "is absolute",
    "Basic data handling: PathIsDir": "is dir",
    "Basic data handling: PathIsFile": "is file",
    "Basic data handling: PathJoin": "join",
    "Basic data handling: PathListDir": "list dir",
    "Basic data handling: PathNormalize": "normalize",
    "Basic data handling: PathOutputDir": "output dir",
    "Basic data handling: PathRelative": "relative",
    "Basic data handling: PathSplit": "split",
    "Basic data handling: PathSplitExt": "splitext",
    "Basic data handling: PathLoadStringFile": "load STRING from file",
    "Basic data handling: PathLoadImageRGB": "load IMAGE from file (RGB)",
    "Basic data handling: PathLoadImageRGBA": "load IMAGE+MASK from file (RGBA)",
    "Basic data handling: PathLoadMaskFromAlpha": "load MASK from alpha channel",
    "Basic data handling: PathLoadMaskFromGreyscale": "load MASK from greyscale/red",
    "Basic data handling: PathSaveStringFile": "save STRING to file",
    "Basic data handling: PathSaveImageRGB": "save IMAGE to file",
    "Basic data handling: PathSaveImageRGBA": "save IMAGE+MASK to file",
}