import { app } from "../../../scripts/app.js";

/*
 * ============================================================================
 * 通用动态端口配置表
 * Key: Python 节点类名
 * Value: prefix (前缀), type (类型), shape (形状)
 * ============================================================================
 */
const DYNAMIC_NODE_CONFIG = {
    // 1. 文本拼接
    "MultipleTextInputs": { 
        prefix: "string_", 
        type: "STRING", 
        shape: 7 
    },
    
    // 2. 遮罩叠加 (RGB)
    "MultipleMasksToImage": { 
        prefix: "mask_",   
        type: "MASK",   
        shape: 7 
    },

    // 3. [新增] 多图拼接
    "MultipleImageStitch": { 
        prefix: "image_", 
        type: "IMAGE", 
        shape: 7 
    },

    // 4. [新增] 多遮罩拼接
    "MultipleMaskStitch": { 
        prefix: "mask_", 
        type: "MASK", 
        shape: 7 
    }
};

/**
 * 将动态端口逻辑应用到指定节点
 */
function applyDynamicPorts(nodeType, config) {
    const originalOnNodeCreated = nodeType.prototype.onNodeCreated || function() {};
    
    nodeType.prototype.onNodeCreated = function () {
        originalOnNodeCreated.apply(this, arguments);

        this.addWidget("button", "Update inputs", null, () => {
            if (!this.inputs) {
                this.inputs = [];
            }
            
            const inputCountWidget = this.widgets.find(w => w.name === "inputcount");
            if (!inputCountWidget) {
                return; 
            }
            
            const target_number_of_inputs = inputCountWidget.value;
            
            // 统计当前符合前缀的输入口
            const num_inputs = this.inputs.filter(input => 
                input.name && input.name.toLowerCase().includes(config.prefix)
            ).length;
            
            if (target_number_of_inputs === num_inputs) return;

            if(target_number_of_inputs < num_inputs){
                // 删
                const inputs_to_remove = num_inputs - target_number_of_inputs;
                for(let i = 0; i < inputs_to_remove; i++) {
                    this.removeInput(this.inputs.length - 1);
                }
            }
            else{
                // 增
                for(let i = num_inputs + 1; i <= target_number_of_inputs; ++i) {
                    const name = `${config.prefix}${i}`;
                    this.addInput(name, config.type, { shape: config.shape });
                }
            }
        });
    }
}

app.registerExtension({
    name: "MyTemplate.DynamicInputs",
    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        if (nodeData.name in DYNAMIC_NODE_CONFIG) {
            const config = DYNAMIC_NODE_CONFIG[nodeData.name];
            applyDynamicPorts(nodeType, config);
        }
    },
});