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 = ("current_filepath", "current_filename", "current_index", "total_file_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 FilePathsEnsureExecution output here to ensure all nodes in the loop execute.
    This ensures all iterations complete before proceeding.
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "current_index": ("INT", {"default": 0, "forceInput": True}),
            },
            "optional": {
                "passthrough": ("*",),
            }
        }
    
    INPUT_IS_LIST = True
    RETURN_TYPES = ("INT",)
    RETURN_NAMES = ("processed_count",)
    FUNCTION = "close_loop"
    CATEGORY = "FilePathsFromFolder/Loop"
    OUTPUT_NODE = True
    
    def close_loop(self, current_index, passthrough=None):
        # 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.
    If you need more than 5 inputs, you can chain multiple of these nodes together.
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "optional": {
                "passthrough_1": ("*",),
                "passthrough_2": ("*",),
                "passthrough_3": ("*",),
                "passthrough_4": ("*",),
                "passthrough_5": ("*",),
            }
        }
    
    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)


class FilePathsLoopSyncTrigger:
    """
    Batch synchronization trigger: splits batch into trigger + batch passthrough.
    - 'trigger' output: single value, connects to nodes that run ONCE (cleanup/checkpoint/etc.)
    - 'anything' output: original batch (N items), connects to FilePathsLoopSyncResume
    Use with FilePathsLoopSyncResume to ensure one-time operation completes before batch continues.
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "anything": ("*", {"forceInput": True}),
            }
        }
    
    INPUT_IS_LIST = True
    RETURN_TYPES = ("*", "*")
    RETURN_NAMES = ("anything", "trigger")
    OUTPUT_IS_LIST = (True, False)
    FUNCTION = "trigger_once"
    CATEGORY = "FilePathsFromFolder/Loop"
    OUTPUT_NODE = False
    
    def trigger_once(self, anything):
        # Split: single trigger for one-time ops + original batch to continue
        trigger_value = anything[0] if isinstance(anything, list) and len(anything) > 0 else anything
        return (anything, trigger_value)


class FilePathsLoopSyncResume:
    """
    Batch synchronization resume: waits for single operation, then continues batch.
    Just a wait point - doesn't modify data.
    Connect batch from Trigger's 'anything' output here.
    Connect completion signal from your one-time operation to 'wait_for'.
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "anything": ("*", {"forceInput": True}),
                "wait_for": ("*",),
            }
        }
    
    INPUT_IS_LIST = True
    RETURN_TYPES = ("*",)
    RETURN_NAMES = ("anything",)
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "resume_batch"
    CATEGORY = "FilePathsFromFolder/Loop"
    OUTPUT_NODE = False
    
    def resume_batch(self, anything, wait_for=None):
        # Simple passthrough after waiting for completion
        return (anything,)


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

NODE_DISPLAY_NAME_MAPPINGS = {
    "LoadFilePathsFromFolder": "Load File Paths From Folder",
    "GetPathByIndex": "Get File Path By Index",
    "FilePathsForLoop": "File Paths For Loop Start",
    "FilePathsLoopClose": "File Paths For Loop End",
    "FilePathsEnsureExecution": "File Paths For Loop Ensure Execution",
    "FilePathsLoopSyncTrigger": "File Paths For Loop Sync Trigger",
    "FilePathsLoopSyncResume": "File Paths For Loop Sync Resume",
}