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

// 通用的動態接點處理函數
const setupSmartDynamicInputs = (nodeType, prefix, type) => {
    const onNodeCreated = nodeType.prototype.onNodeCreated;
    nodeType.prototype.onNodeCreated = function () {
        onNodeCreated?.apply(this, arguments);
        const inputCountWidget = this.widgets.find(w => w.name === "input_count");

        inputCountWidget.callback = () => {
            const targetCount = inputCountWidget.value;
            
            // 1. 找出目前所有屬於該 prefix 的輸入點
            const currentInputs = this.inputs.filter(i => i.name.startsWith(`${prefix}_`));
            const currentCount = currentInputs.length;

            if (targetCount > currentCount) {
                // 2. 如果目標數量較多：只新增缺少的接點，原有的接點不受影響
                for (let i = currentCount; i < targetCount; i++) {
                    this.addInput(`${prefix}_${i}`, type);
                }
            } else if (targetCount < currentCount) {
                // 3. 如果目標數量較少：從最後面開始移除，保留前面的連線
                for (let i = currentCount - 1; i >= targetCount; i--) {
                    // 找到該名稱對應的索引並移除
                    const inputIdx = this.inputs.findIndex(input => input.name === `${prefix}_${i}`);
                    if (inputIdx !== -1) {
                        this.removeInput(inputIdx);
                    }
                }
            }

            // 重新整理 UI 配置
            this.setSize(this.computeSize());
        };

        // 延遲初始化，確保節點建立時 UI 正確
        setTimeout(() => {
            if (inputCountWidget.callback) inputCountWidget.callback();
        }, 50);
    };
};

app.registerExtension({
    name: "CustomNodes.IterSwitches",
    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        // 根據節點名稱套用不同的前綴與類型
        if (nodeData.name === "Iter_TextSwitch") {
            setupSmartDynamicInputs(nodeType, "text", "STRING");
        } else if (nodeData.name === "Iter_ConditioningSwitch") {
            setupSmartDynamicInputs(nodeType, "cond", "CONDITIONING");
        } else if (nodeData.name === "Iter_CLIPSwitch") {
            setupSmartDynamicInputs(nodeType, "clip", "CLIP");
        } else if (nodeData.name === "Iter_ImageSwitch") {
            setupSmartDynamicInputs(nodeType, "image", "IMAGE");
        } else if (nodeData.name === "Iter_ModelSwitch") {
            setupSmartDynamicInputs(nodeType, "model", "MODEL");
        } else if (nodeData.name === "Iter_FloatSwitch") {
            setupSmartDynamicInputs(nodeType, "float", "FLOAT");
        }
    }
});