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

// Bonso Switch Model: dos sockets MODEL y una casilla `switch`.
//
// ComfyUI valida el prompt COMPLETO antes de ejecutar (recorre todos los
// enlaces desde los nodos de salida), sin preguntar por el `switch`. Si la
// rama que NO toca esta rota (nodo en bypass, entrada que falta), el
// trabajo muere en la validacion aunque jamas se vaya a usar. Para que la
// rama inactiva se ignore de verdad hay que quitarla del prompt, y eso no
// lo puede hacer el servidor: lo hace aqui, envolviendo `graphToPrompt`.
//
// Solo se poda cuando se puede saber el valor del `switch`:
//   - casilla literal en el propio nodo, o
//   - enlace desde `Bonso Booleano` (su `estado` viaja serializado en su
//     propio prompt, que este nodo mira).
// En cualquier otro caso no se poda nada y vale el comportamiento de siempre.
// El grafo del usuario nunca se toca: el nodo conserva sus dos lineas.

app.registerExtension({
    name: "BonsoTools.SwitchModel",

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

        const original = app.graphToPrompt;
        app.graphToPrompt = async function () {
            const r = await original.apply(this, arguments);
            try {
                podarRamaInactiva(r && r.output);
            } catch (e) {
                console.warn("[Bonso Switch Model] no se pudo podar la rama inactiva", e);
            }
            return r;
        };
    },
});

function valorDeSwitch(output, entrada) {
    if (typeof entrada === "boolean") return entrada;
    if (typeof entrada === "number") return entrada !== 0;
    if (Array.isArray(entrada)) {
        const origen = output[String(entrada[0])];
        if (origen && origen.class_type === "Bonso_Booleano") {
            const estado = origen.inputs ? origen.inputs.estado : undefined;
            if (typeof estado === "boolean") return estado;
            if (typeof estado === "number") return estado !== 0;
        }
    }
    return null;
}

function podarRamaInactiva(output) {
    if (!output || typeof output !== "object") return 0;
    let quitadas = 0;
    for (const id of Object.keys(output)) {
        const nodo = output[id];
        if (!nodo || nodo.class_type !== "Bonso_Switch_Model") continue;
        const inputs = nodo.inputs || (nodo.inputs = {});
        const activo = valorDeSwitch(output, inputs.switch);
        if (activo === null) continue;
        const rama = activo ? "false_model" : "true_model";
        if (rama in inputs) {
            delete inputs[rama];
            quitadas++;
        }
    }
    if (quitadas) {
        console.debug(`[Bonso Switch Model] rama inactiva fuera del prompt (${quitadas})`);
    }
    return quitadas;
}