import os
import random
import folder_paths

class CheckpointSelector:
    def __init__(self):
        self.current_index = 0
        self.current_count = 0
        self.last_folder = None
        self.last_selection = None
        self.last_manual_index = 0

    @classmethod
    def INPUT_TYPES(s):
        # Find folders in 'checkpoints' paths
        checkpoint_paths = folder_paths.get_folder_paths("checkpoints")
        checkpoint_folders = []
        
        for path in checkpoint_paths:
            if os.path.exists(path) and os.path.isdir(path):
                # List subdirectories
                try:
                    subdirs = [d for d in os.listdir(path) if os.path.isdir(os.path.join(path, d))]
                    checkpoint_folders.extend(subdirs)
                except Exception:
                    pass
        
        # Remove duplicates and sort
        checkpoint_folders = sorted(list(set(checkpoint_folders)))
        
        if not checkpoint_folders:
            checkpoint_folders = ["None"]

        return {
            "required": {
                "folder": (checkpoint_folders, ),
                "mode": (["Random", "Sequential", "Manual"], {"default": "Random"}),
                "repeat_count": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
                "manual_index": ("INT", {"default": 0, "min": 0, "max": 10000, "step": 1}),
            }
        }

    RETURN_TYPES = ("STRING", "INT")
    RETURN_NAMES = ("checkpoint_name", "total_checkpoints")
    FUNCTION = "select_checkpoint"
    CATEGORY = "AODH Pack"

    def select_checkpoint(self, folder, mode, repeat_count, manual_index):
        repeat_count = max(1, repeat_count)
        if folder == "None":
            return ("", 0)

        # Find the full path for the selected folder
        checkpoint_paths = folder_paths.get_folder_paths("checkpoints")
        target_path = None
        
        for path in checkpoint_paths:
            check_path = os.path.join(path, folder)
            if os.path.exists(check_path) and os.path.isdir(check_path):
                target_path = check_path
                break
        
        if not target_path:
            return ("", 0)

        # Find checkpoint files (.safetensors, .ckpt)
        extensions = ('.safetensors', '.ckpt')
        checkpoint_files = [f for f in os.listdir(target_path) if f.lower().endswith(extensions)]
        checkpoint_files.sort()
        
        count = len(checkpoint_files)
        if count == 0:
            return ("", 0)

        selected_checkpoint = ""
        
        # Reset state if folder changes
        if self.last_folder != folder:
            self.current_index = 0
            self.current_count = 0
            self.last_selection = None
            self.last_folder = folder

        # Reset state if manual_index changes (for Sequential mode)
        if self.last_manual_index != manual_index:
            self.current_index = 0
            self.current_count = 0
            self.last_manual_index = manual_index

        if mode == "Random":
            if self.current_count >= repeat_count or self.last_selection is None:
                rng = random.Random()
                self.last_selection = rng.choice(checkpoint_files)
                self.current_count = 0
            
            selected_checkpoint = self.last_selection
            self.current_count += 1
        
        elif mode == "Sequential":
            if self.current_count >= repeat_count:
                self.current_index += 1
                self.current_count = 0
            
            selected_checkpoint = checkpoint_files[(self.current_index + manual_index) % count]
            self.current_count += 1

        elif mode == "Manual":
            selected_checkpoint = checkpoint_files[manual_index % count]

        # Construct the relative path for ComfyUI Checkpoint loader
        # Use forward slashes for consistency with ComfyUI paths
        full_checkpoint_name = f"{folder}/{selected_checkpoint}"
        
        return (full_checkpoint_name, count)

    @classmethod
    def IS_CHANGED(s, folder, mode, repeat_count, manual_index):
        if mode == "Random" or mode == "Sequential":
            return float("nan")
        return f"{folder}_{manual_index}"