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

const DEFAULT_SLOTS = 2;

function serializeSlots(node) {
    const entries = (node._loraSlots ?? []).map(s => ({
        name:           s.nameW.value,
        strength_model: s.smW.value,
        enabled:        s.enW.value,
    }));
    const jw = node.widgets?.find(w => w.name === "loras_json");
    if (jw) jw.value = JSON.stringify(entries);
}

function addSlot(node, def = {}) {
    const loras = node._loraNames ?? ["None"];

    const nameW = node.addWidget("combo", "lora", def.name ?? "None",
        () => serializeSlots(node), { values: loras });

    const smW = node.addWidget("number", "strength", def.strength_model ?? 1.0,
        () => serializeSlots(node), { min: 0.0, max: 2.0, step: 0.1, precision: 2 });

    const enW = node.addWidget("toggle", "enabled", def.enabled ?? true,
        () => serializeSlots(node));

    node._loraSlots = node._loraSlots ?? [];
    node._loraSlots.push({ nameW, smW, enW });

    serializeSlots(node);
    node.setSize([node.size[0], node.computeSize()[1]]);
    node.graph?.setDirtyCanvas(true, true);
}

function removeSlot(node) {
    if (!node._loraSlots?.length) return;
    const s = node._loraSlots.pop();
    for (const w of [s.nameW, s.smW, s.enW]) {
        const i = node.widgets.indexOf(w);
        if (i !== -1) node.widgets.splice(i, 1);
    }
    serializeSlots(node);
    node.setSize([node.size[0], node.computeSize()[1]]);
    node.graph?.setDirtyCanvas(true, true);
}

function init(node) {
    if (node._atomSimpleInit) return;
    node._atomSimpleInit = true;
    node._loraSlots      = [];

    const jw = node.widgets?.find(w => w.name === "loras_json");
    if (jw) {
        jw.hidden = true;
        jw.computeSize = () => [0, -4];
    }

    node.addWidget("button", "+ Add LoRA",    null, () => addSlot(node));
    node.addWidget("button", "− Remove LoRA", null, () => removeSlot(node));

    let restored = false;
    if (jw?.value && jw.value !== "[]") {
        try {
            const entries = JSON.parse(jw.value);
            if (entries.length > 0) {
                entries.forEach(e => addSlot(node, e));
                restored = true;
            }
        } catch {}
    }
    if (!restored) {
        for (let i = 0; i < DEFAULT_SLOTS; i++) addSlot(node);
    }
}

app.registerExtension({
    name: "Atom_AI.SimpleLoRANode",

    nodeCreated(node) {
        if (node.comfyClass !== "AtomSimpleLoRANode") return;
        node._loraNames = app._atomLoraNames ?? ["None"];
        init(node);
    },
});