import os
import glob


class FilePathsForLoop:
    """
    Iterates through files in a folder one at a time.
    Returns each file path individually for processing in a loop.
    Connect to any node that processes files, then to FilePathsLoopClose.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "folder_path": ("STRING", {"default": ""}),
                "extension": ("STRING", {"default": "*.png"}),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "INT", "INT")
    RETURN_NAMES = ("filepath", "filename", "current_index", "total_count")
    FUNCTION = "iterate_files"
    CATEGORY = "FilePathsFromFolder/Loop"
    OUTPUT_NODE = False
    OUTPUT_IS_LIST = (True, True, True, True)
    
    def iterate_files(self, folder_path, extension):
        # Build search pattern
        search_pattern = os.path.join(folder_path, extension)
        
        # Find all matching files
        matched_files = glob.glob(search_pattern)
        matched_files.sort()
        
        total_count = len(matched_files)
        
        # If no files, return empty lists
        if not matched_files:
            return ([], [], [], [])
        
        # Build output lists
        filepaths = []
        filenames = []
        indices = []
        counts = []
        
        for idx, filepath in enumerate(matched_files):
            basename = os.path.basename(filepath)
            filename_without_ext = os.path.splitext(basename)[0]
            
            filepaths.append(filepath)
            filenames.append(filename_without_ext)
            indices.append(idx)
            counts.append(total_count)
        
        # Return as lists - OUTPUT_IS_LIST tells ComfyUI this is a batch
        return (filepaths, filenames, indices, counts)


class FilePathsLoopClose:
    """
    Closes the file paths loop.
    Connect any outputs you want to collect from the loop iterations.
    This ensures all iterations complete before proceeding.
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "current_index": ("INT", {"default": 0, "forceInput": True}),
            },
            "optional": {
                "passthrough_1": ("*",),
                "passthrough_2": ("*",),
                "passthrough_3": ("*",),
                "passthrough_4": ("*",),
                "passthrough_5": ("*",),
            }
        }
    
    INPUT_IS_LIST = True
    RETURN_TYPES = ("INT",)
    RETURN_NAMES = ("total_iterations",)
    FUNCTION = "close_loop"
    CATEGORY = "FilePathsFromFolder/Loop"
    OUTPUT_NODE = True
    
    def close_loop(self, current_index, **kwargs):
        # Count total iterations
        total = len(current_index) if isinstance(current_index, list) else 1
        return (total,)


class FilePathsEnsureExecution:
    """
    Ensures connected nodes execute by passing through their outputs.
    Use this to chain multiple nodes that need to execute during each loop iteration.
    Connect up to 10 node outputs, then connect this node's output to FilePathsLoopClose.
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "optional": {
                "input_1": ("*",),
                "input_2": ("*",),
                "input_3": ("*",),
                "input_4": ("*",),
                "input_5": ("*",),
                "input_6": ("*",),
                "input_7": ("*",),
                "input_8": ("*",),
                "input_9": ("*",),
                "input_10": ("*",),
            }
        }
    
    RETURN_TYPES = ("*",)
    RETURN_NAMES = ("passthrough",)
    FUNCTION = "passthrough"
    CATEGORY = "FilePathsFromFolder/Loop"
    OUTPUT_NODE = False
    
    def passthrough(self, **kwargs):
        # Return the first non-None input as passthrough
        for key in sorted(kwargs.keys()):
            if kwargs[key] is not None:
                return (kwargs[key],)
        return (None,)


class LoadFilePathsFromFolder:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "folder_path": ("STRING", {"default": ""}),
                "extension": ("STRING", {"default": "*.png"}),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "INT")
    RETURN_NAMES = ("filepaths", "filenames", "count")
    FUNCTION = "load_files"
    CATEGORY = "FilePathsFromFolder"
    
    def load_files(self, folder_path, extension):
        # Build search pattern
        search_pattern = os.path.join(folder_path, extension)
        
        # Find all matching files
        matched_files = glob.glob(search_pattern)
        matched_files.sort()
        
        # Extract file paths and names
        filepaths = []
        filenames = []
        
        for filepath in matched_files:
            filepaths.append(filepath)
            # Get filename without extension
            basename = os.path.basename(filepath)
            filename_without_ext = os.path.splitext(basename)[0]
            filenames.append(filename_without_ext)
        
        # Join with line breaks
        filepaths_str = "\n".join(filepaths)
        filenames_str = "\n".join(filenames)
        count = len(matched_files)
        
        return (filepaths_str, filenames_str, count)


class GetPathByIndex:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "filepaths": ("STRING", {"default": "", "multiline": True}),
                "filenames": ("STRING", {"default": "", "multiline": True}),
                "index": ("INT", {"default": 0, "min": 0}),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("filepath", "filename")
    FUNCTION = "get_by_index"
    CATEGORY = "FilePathsFromFolder"
    
    def get_by_index(self, filepaths, filenames, index):
        # Split by line breaks
        filepath_list = filepaths.split("\n") if filepaths else []
        filename_list = filenames.split("\n") if filenames else []
        
        # Get item at index (default to empty if out of range)
        filepath = filepath_list[index] if index < len(filepath_list) else ""
        filename = filename_list[index] if index < len(filename_list) else ""
        
        return (filepath, filename)


NODE_CLASS_MAPPINGS = {
    "LoadFilePathsFromFolder": LoadFilePathsFromFolder,
    "GetPathByIndex": GetPathByIndex,
    "FilePathsForLoop": FilePathsForLoop,
    "FilePathsLoopClose": FilePathsLoopClose,
    "FilePathsEnsureExecution": FilePathsEnsureExecution,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LoadFilePathsFromFolder": "Load File Paths From Folder",
    "GetPathByIndex": "Get Path By Index",
    "FilePathsForLoop": "File Paths For Loop",
    "FilePathsLoopClose": "File Paths Loop Close",
    "FilePathsEnsureExecution": "File Paths Ensure Execution",
}