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

async function getAvailableLoras() {
    try {
        const response = await api.fetchApi("/object_info");
        const objectInfo = await response.json();
        for (const [, nodeInfo] of Object.entries(objectInfo)) {
            if (nodeInfo.input?.required?.lora_name) {
                return nodeInfo.input.required.lora_name[0] || ["None"];
            }
        }
        return ["None"];
    } catch (error) {
        console.error("Error fetching LoRAs:", error);
        return ["None"];
    }
}

app.registerExtension({
    name: "MultiLoRAStack",

    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        if (nodeData.name !== "MultiLoRAStack") return;

        const onNodeCreated = nodeType.prototype.onNodeCreated;
        nodeType.prototype.onNodeCreated = async function () {
            const result = onNodeCreated?.apply(this, arguments);

            const stackWidget = this.widgets?.find(w => w.name === "lora_stack");
            if (stackWidget) {
                stackWidget.computeSize = () => [0, -4];
                stackWidget.draw = () => {};
            }

            this.size = [450, 300];
            this.loraData = [];
            this.availableLoras = await getAvailableLoras();

            if (!this._configureWasRun) {
                for (let i = (this.inputs?.length ?? 0) - 1; i >= 0; i--) {
                    if (/^strength_\d+$/.test(this.inputs[i].name)) this.removeInput(i);
                }
            }

            this.addWidget("button", "➕ Add LoRA", null, () => { this.addLoRA(); });
            this.addWidget("button", "Toggle All LoRAs", null, () => { this.toggleAllLoRAs(); });

            this.addLoRA();
            return result;
        };

        nodeType.prototype.addLoRA = function(initialData = null) {
            const index    = this.loraData.length;
            const on       = initialData?.on       ?? true;
            const lora     = initialData?.lora     ?? "None";
            const strength = initialData?.strength ?? 1.0;

            this.loraData.push({ on, lora, strength });

            const enableWidget = this.addWidget("toggle", `Enable LoRA ${index + 1}`, on, (value) => {
                this.loraData[index].on = value;
                this.updateLoraStack();
            });

            const loraWidget = this.addWidget("combo", `LoRA ${index + 1}`, lora, (value) => {
                this.loraData[index].lora = value;
                this.updateLoraStack();
            }, { values: this.availableLoras });

            const strengthWidget = this.addWidget("number", `Strength ${index + 1}`, strength, (value) => {
                this.loraData[index].strength = value;
                this.updateLoraStack();
            }, { min: -2.0, max: 2.0, step: 0.1 });

            const removeWidget = this.addWidget("button", `Remove LoRA ${index + 1}`, null, () => {
                this.removeLoRA(index);
            });

            this.loraData[index]._widgets = {
                enable: enableWidget,
                lora: loraWidget,
                strength: strengthWidget,
                remove: removeWidget
            };

            const slotName = `strength_${index + 1}`;
            const existing = this.inputs?.find(inp => inp.name === slotName);
            if (existing) {
                if (!existing.widget) existing.widget = { name: strengthWidget.name };
            } else {
                this.addInput(slotName, "FLOAT", { widget: { name: strengthWidget.name } });
            }

            this.updateLoraStack();
            this.setSize([this.size[0], this.computeSize()[1]]);
        };

        const onConfigure = nodeType.prototype.onConfigure;
        nodeType.prototype.onConfigure = function(info) {
            this._configureWasRun = true;
            const result = onConfigure?.apply(this, arguments);

            const stackWidget = this.widgets?.find(w => w.name === "lora_stack");
            const savedJson = stackWidget?.value;

            const tryRestore = () => {
                if (!this.availableLoras) { setTimeout(tryRestore, 50); return; }
                if (!savedJson) return;
                let saved;
                try { saved = JSON.parse(savedJson); } catch (e) { return; }
                if (!Array.isArray(saved) || saved.length === 0) return;

                for (const item of this.loraData) {
                    if (item._widgets) {
                        Object.values(item._widgets).forEach(w => {
                            const i = this.widgets.indexOf(w);
                            if (i !== -1) this.widgets.splice(i, 1);
                        });
                    }
                }
                this.loraData = [];

                for (const item of saved) {
                    this.addLoRA(item);
                }
                this.setSize([this.size[0], this.computeSize()[1]]);
            };
            setTimeout(tryRestore, 0);
            return result;
        };

        nodeType.prototype.removeLoRA = function(index) {
            const loraItem = this.loraData[index];
            if (loraItem?._widgets) {
                Object.values(loraItem._widgets).forEach(widget => {
                    const widgetIndex = this.widgets.indexOf(widget);
                    if (widgetIndex !== -1) this.widgets.splice(widgetIndex, 1);
                });
            }

            const slotName = `strength_${index + 1}`;
            const inputIdx = this.inputs?.findIndex(inp => inp.name === slotName) ?? -1;
            if (inputIdx !== -1) this.removeInput(inputIdx);

            this.loraData.splice(index, 1);
            this.updateWidgetNames();
            this.updateLoraStack();
            this.setSize([this.size[0], this.computeSize()[1]]);
        };

        nodeType.prototype.updateWidgetNames = function() {
            this.loraData.forEach((loraItem, index) => {
                if (loraItem._widgets) {
                    loraItem._widgets.enable.name   = `Enable LoRA ${index + 1}`;
                    loraItem._widgets.lora.name     = `LoRA ${index + 1}`;
                    loraItem._widgets.strength.name = `Strength ${index + 1}`;
                    loraItem._widgets.remove.name   = `Remove LoRA ${index + 1}`;
                }
            });
            let n = 1;
            for (const inp of (this.inputs || [])) {
                if (/^strength_\d+$/.test(inp.name)) {
                    inp.name   = `strength_${n}`;
                    inp.widget = { name: `Strength ${n}` };
                    n++;
                }
            }
        };

        nodeType.prototype.toggleAllLoRAs = function() {
            const allEnabled = this.loraData.every(lora => lora.on);
            const newState = !allEnabled;
            this.loraData.forEach((loraItem) => {
                loraItem.on = newState;
                if (loraItem._widgets?.enable) loraItem._widgets.enable.value = newState;
            });
            this.updateLoraStack();
        };

        nodeType.prototype.updateLoraStack = function() {
            const stackWidget = this.widgets?.find(w => w.name === "lora_stack");
            if (stackWidget) {
                stackWidget.value = JSON.stringify(
                    this.loraData.map(item => ({ on: item.on, lora: item.lora, strength: item.strength }))
                );
            }
        };

        const originalComputeSize = nodeType.prototype.computeSize;
        nodeType.prototype.computeSize = function(out) {
            const size = originalComputeSize ? originalComputeSize.call(this, out) : [300, 200];
            size[0] = Math.max(size[0], 450);
            return size;
        };
    }
});