# multi_controlnet_loader.py
import os
import folder_paths
from comfy import model_management
import comfy.controlnet
import comfy.utils

class MultiControlNetLoader:
    @classmethod
    def INPUT_TYPES(cls):
        # ControlNetモデルリストを取得
        controlnet_list = folder_paths.get_filename_list("controlnet")
        
        return {
            "required": {
                "checkpoint": ("MODEL", {"default": None}),
                "sd_cn_name": (controlnet_list, {"default": controlnet_list[0] if controlnet_list else ""}),
                "xl_cn_name": (controlnet_list, {"default": controlnet_list[0] if controlnet_list else ""}),
            }
        }

    RETURN_TYPES = ("CONTROL_NET",)
    FUNCTION = "switch_and_load"
    CATEGORY = "ControlNet"

    def switch_and_load(self, checkpoint, sd_cn_name, xl_cn_name):
        """
        checkpoint 名に "sdxl" を含む場合は SDXL 用モデルを、それ以外は SD 用モデルをロードします。
        """
        try:
            # checkpoint 名で SD or SDXL を判定
            # checkpointがModelPatcherオブジェクトの場合の処理
            if hasattr(checkpoint, 'model') and hasattr(checkpoint.model, 'model_config'):
                model_config = checkpoint.model.model_config
                # SDXL判定 - model_configからSDXLかどうかを判定
                is_sdxl = (
                    hasattr(model_config, 'unet_config') and 
                    model_config.unet_config.get('model_channels', 320) == 320 and
                    model_config.unet_config.get('adm_in_channels', None) is not None
                ) or str(model_config).lower().find('sdxl') != -1
            else:
                # フォールバック: 文字列として名前をチェック
                name = str(checkpoint).lower()
                is_sdxl = "sdxl" in name
            
            # 使用するControlNetモデル名を決定
            target_name = xl_cn_name if is_sdxl else sd_cn_name
            
            # ControlNetモデルのパスを取得
            controlnet_path = folder_paths.get_full_path("controlnet", target_name)
            if controlnet_path is None:
                raise FileNotFoundError(f"ControlNet model '{target_name}' not found in controlnet folders")
            
            # ControlNetモデルをロード
            controlnet = comfy.controlnet.load_controlnet(controlnet_path)
            
            print(f"MultiControlNetLoader: Loaded {'SDXL' if is_sdxl else 'SD'} ControlNet model: {target_name}")
            
            return (controlnet,)
            
        except Exception as e:
            print(f"Error in MultiControlNetLoader: {str(e)}")
            print(f"Attempted to load: {'SDXL' if 'is_sdxl' in locals() and is_sdxl else 'SD'} model with name: {target_name if 'target_name' in locals() else 'Unknown'}")
            raise e

# ノードクラスマッピング
NODE_CLASS_MAPPINGS = {"MultiControlNetLoader": MultiControlNetLoader}
NODE_DISPLAY_NAME_MAPPINGS = {"MultiControlNetLoader": "Multi ControlNet Loader"}