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

const DEFAULT_SLOTS = 1;

function serializeSlots(node) {
    const entries = (node._loraSlots ?? []).map(s => ({
        name:           s.nameW.value,
        strength_model: s.smW.value,
        strength_clip:  s.scW.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", "model str.", def.strength_model ?? 1.0,
        () => serializeSlots(node), { min: -10, max: 10, step: 0.1, precision: 2 });

    const scW = node.addWidget("number", "clip str.", def.strength_clip ?? 1.0,
        () => serializeSlots(node), { min: -10, max: 10, 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, scW, 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.scW, 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._atomDynInit) return;
    node._atomDynInit = true;
    node._loraSlots   = [];

    // Hide loras_json
    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));

    // Restore from saved JSON or add defaults
    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);
    }
}

// Fetch lora list directly from ComfyUI's /object_info endpoint
async function fetchLoraNames() {
    try {
        const r = await fetch("/object_info");
        const all = await r.json();

        // Try every known lora loader node name
        const candidates = [
            "LoraLoader",
            "LoRALoader", 
            "Load LoRA",
            "LoraLoaderModelOnly",
        ];

        for (const name of candidates) {
            const def = all[name];
            if (!def) continue;
            // Search all required inputs for a list that looks like lora filenames
            for (const [, val] of Object.entries(def.input?.required ?? {})) {
                if (Array.isArray(val[0]) && val[0].some(v => 
                    typeof v === "string" && (v.endsWith(".safetensors") || v.endsWith(".pt") || v.endsWith(".ckpt"))
                )) {
                    return ["None", ...val[0]];
                }
            }
        }

        // Last resort: scan ALL node defs for a safetensors list
        for (const def of Object.values(all)) {
            for (const [, val] of Object.entries(def.input?.required ?? {})) {
                if (Array.isArray(val[0]) && val[0].length > 0 &&
                    val[0].some(v => typeof v === "string" && v.includes("lora"))) {
                    return ["None", ...val[0]];
                }
            }
        }
    } catch (e) {
        console.warn("[Atom_AI] Could not fetch lora list:", e);
    }
    return ["None"];
}

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

    async setup() {
        // Fetch and cache lora list at startup
        app._atomLoraNames = await fetchLoraNames();
        console.log("[Atom_AI] LoRA list loaded:", app._atomLoraNames.length, "entries");
    },

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