"""
Knishika Dual Lora Stack - ComfyUI Custom Node
Author: knishika
Description: Load multiple Lora models to two separate model paths with shared CLIP
"""

import folder_paths

class KnishikaDualLoraStack:
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_1": ("MODEL",),
                "model_2": ("MODEL",),
                "clip": ("CLIP",),
                "lora_01": (["None"] + folder_paths.get_filename_list("loras"),),
                "lora_01_strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "lora_02": (["None"] + folder_paths.get_filename_list("loras"),),
                "lora_02_strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "lora_03": (["None"] + folder_paths.get_filename_list("loras"),),
                "lora_03_strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "lora_04": (["None"] + folder_paths.get_filename_list("loras"),),
                "lora_04_strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "lora_05": (["None"] + folder_paths.get_filename_list("loras"),),
                "lora_05_strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                "lora_06": (["None"] + folder_paths.get_filename_list("loras"),),
                "lora_06_strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
            }
        }

    RETURN_TYPES = ("MODEL", "MODEL", "CLIP")
    RETURN_NAMES = ("model_1", "model_2", "clip")
    FUNCTION = "load_lora"
    CATEGORY = "knishika/loaders"

    def load_lora(self, model_1, model_2, clip, lora_01, lora_01_strength, lora_02, lora_02_strength, lora_03, lora_03_strength, lora_04, lora_04_strength, lora_05, lora_05_strength, lora_06, lora_06_strength):
        from nodes import LoraLoader
        
        lora_loader = LoraLoader()
        
        # Process each lora for both models
        loras = [
            (lora_01, lora_01_strength),
            (lora_02, lora_02_strength),
            (lora_03, lora_03_strength),
            (lora_04, lora_04_strength),
            (lora_05, lora_05_strength),
            (lora_06, lora_06_strength)
        ]
        
        current_model_1 = model_1
        current_model_2 = model_2
        current_clip = clip
        
        for lora_name, strength in loras:
            if lora_name != "None" and strength != 0:
                # Apply lora to model_1
                current_model_1, current_clip = lora_loader.load_lora(current_model_1, current_clip, lora_name, strength, strength)
                
                # Apply the same lora to model_2 (create fresh clip reference for model_2)
                current_model_2, _ = lora_loader.load_lora(current_model_2, current_clip, lora_name, strength, strength)
        
        return (current_model_1, current_model_2, current_clip)

NODE_CLASS_MAPPINGS = {
    "KnishikaDualLoraStack": KnishikaDualLoraStack,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "KnishikaDualLoraStack": "Knishika Dual Lora Stack",
}