import os
import random
import folder_paths

class LoraSelector:
    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 'Illustrious' folders in all lora paths
        lora_paths = folder_paths.get_folder_paths("loras")
        illustrious_folders = []
        
        for path in lora_paths:
            illustrious_path = os.path.join(path, "Illustrious")
            if os.path.exists(illustrious_path) and os.path.isdir(illustrious_path):
                # List subdirectories
                try:
                    subdirs = [d for d in os.listdir(illustrious_path) if os.path.isdir(os.path.join(illustrious_path, d))]
                    illustrious_folders.extend(subdirs)
                except Exception:
                    pass
        
        # Remove duplicates and sort
        illustrious_folders = sorted(list(set(illustrious_folders)))
        
        if not illustrious_folders:
            illustrious_folders = ["None"]

        return {
            "required": {
                "folder": (illustrious_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 = ("lora_name", "total_loras")
    FUNCTION = "select_lora"
    CATEGORY = "AODH Pack"

    def select_lora(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
        lora_paths = folder_paths.get_folder_paths("loras")
        target_path = None
        
        for path in lora_paths:
            check_path = os.path.join(path, "Illustrious", 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 .safetensors files
        lora_files = [f for f in os.listdir(target_path) if f.endswith('.safetensors')]
        lora_files.sort() # Ensure consistent order
        
        count = len(lora_files)
        if count == 0:
            return ("", 0)

        selected_lora = ""
        
        # 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:
                # Use a local Random instance to ensure randomness regardless of global seed
                rng = random.Random()
                self.last_selection = rng.choice(lora_files)
                self.current_count = 0
            
            selected_lora = 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_lora = lora_files[(self.current_index + manual_index) % count]
            self.current_count += 1

        elif mode == "Manual":
            selected_lora = lora_files[manual_index % count]

        # Construct the relative path for ComfyUI Lora loader
        # Use forward slashes for consistency
        full_lora_name = f"Illustrious/{folder}/{selected_lora}"
        
        return (full_lora_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}"