import os
import numpy as np
import torch
from PIL import Image, ImageOps

class LoadLoopImagesFromDir:
    last_index = 0  # Class attribute to track the last accessed image

    @classmethod
    def INPUT_TYPES(cls):
        """
        The function `INPUT_TYPES` defines required and optional input types with default values and
        labels for a specified class.
        
        :param cls: The `INPUT_TYPES` function defines input parameters for a class. The parameters are
        categorized as "required" and "optional" inputs
        :return: A dictionary containing information about input types for a class. The dictionary has
        two keys: "required" and "optional", each containing a dictionary of input types with their
        corresponding data types and optional parameters such as default values and labels.
        """
        return {
            "required": {
                "directory": ("STRING", {"default": ""}),
            },
            "optional": {
                "load_always": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}),
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK", "STRING")
    RETURN_NAMES = ("IMAGE", "MASK", "FILE PATH")
    OUTPUT_IS_LIST = (False, False, False)  # Only return one image at a time

    FUNCTION = "load_images"

    CATEGORY = "image"

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        """
        The function `IS_CHANGED` increments `last_index` every time it is called, looping when needed,
        and returns a hash based on the index and input arguments.
        
        :param cls: The `cls` parameter in the given function represents the class itself. It is used to
        access class variables and methods within the function
        :return: The function `IS_CHANGED` returns the hash value of a tuple containing `cls.last_index`
        and a frozenset of the keyword arguments `kwargs.items()`.
        """
        if kwargs.get("load_always", False):
            return float("NaN")

        directory = kwargs.get("directory", "")
        if not os.path.isdir(directory):
            return hash(frozenset(kwargs.items()))

        valid_extensions = {'.jpg', '.jpeg', '.png', '.webp'}
        dir_files = sorted(
            f for f in os.listdir(directory)
            if any(f.lower().endswith(ext) for ext in valid_extensions)
        )

        if dir_files:
            cls.last_index = (cls.last_index + 1) % len(dir_files)

        return hash((cls.last_index, frozenset(kwargs.items())))

    def load_images(cls, directory: str, load_always=False):
        """
        The function `load_images` loads image files from a specified directory, processes them, and
        returns the 
        
        image, mask (if alpha channel exists), and the image path.
        
        :param directory: The `directory` parameter in the `load_images` method is a string that
        represents the path to the directory containing the image files that you want to load. This
        method loads images from the specified directory based on the valid image file extensions (.jpg,
        .jpeg, .png, .webp)
        :type directory: str
        :param load_always: The `load_always` parameter in the `load_images` method is a boolean flag
        that determines whether to always load images from the specified directory or not. By default,
        it is set to `False`, which means that images will be loaded cyclically from the directory based
        on the `, defaults to False (optional)
        :return: The `load_images` method returns three values: `image`, `mask`, and `image_path`.
        """
        if not os.path.isdir(directory):
            raise FileNotFoundError(f"Directory '{directory}' not found.")

        # Get all valid image files
        valid_extensions = {'.jpg', '.jpeg', '.png', '.webp'}
        dir_files = sorted(
            f for f in os.listdir(directory)
            if any(f.lower().endswith(ext) for ext in valid_extensions)
        )

        if not dir_files:
            raise FileNotFoundError(f"No image files found in directory '{directory}'.")

        # Cycle through images using `last_index`
        image_count = len(dir_files)
        image_path = os.path.join(directory, dir_files[cls.last_index])

        # Update index for the next call (loop back if at the end)
        cls.last_index = (cls.last_index + 1) % image_count

        # Load image
        img = Image.open(image_path)
        img = ImageOps.exif_transpose(img)
        image = img.convert("RGB")
        image = np.array(image).astype(np.float32) / 255.0
        image = torch.from_numpy(image)[None,]  # Add batch dimension

        # Load mask (if alpha channel exists)
        if 'A' in img.getbands():
            mask = np.array(img.getchannel('A')).astype(np.float32) / 255.0
            mask = 1.0 - torch.from_numpy(mask)
        else:
            mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")

        return image, mask, str(image_path)

# Dictionary to map node classes to their names
NODE_CLASS_MAPPINGS = {
    "LoadLoopImagesFromDir": LoadLoopImagesFromDir,
}

# Dictionary to map node names to their friendly display names
NODE_DISPLAY_NAME_MAPPINGS = {
    "LoadLoopImagesFromDir": "Load Loop Images From Dir",
}