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

// Bonso Switch: dos sockets `*` (cualquier tipo) y una casilla `switch`.
//
// Trabajo 1 - Poda de la rama inactiva (mismo truco que
// js/bonso_switch_model.js): ComfyUI valida el prompt COMPLETO antes de
// ejecutar, asi que un error de validacion en la rama que NO toca (nodo en
// bypass con entrada requerida vacia, tipo mal enchufado) reventaria el
// encolado aunque esa rama nunca se fuera a usar. El backend no puede
// podar: se envuelve `graphToPrompt` y se borra la rama inactiva cuando el
// valor de `switch` se puede saber (casilla literal o enlace desde
// `Bonso Booleano`). En cualquier otro caso no se poda nada.
//
// Trabajo 2 - Restriccion de tipo en el lienzo: `onConnectInput` devuelve
// false si el tipo que llega no coincide con el tipo ya enchufado en la
// otra entrada de datos (`on true` / `on false`; la primera conexion manda;
// `*` no restringe nada). La casilla `switch` NO participa: es BOOLEAN y
// acepta el booleano de Bonso Booleano sin mas. Es restriccion VISUAL: el
// servidor no la conoce porque la salida es `*`.

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

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

        // --- Trabajo 1: poda al encolar (una sola envoltura global) ---
        if (!app.__bonsoSwitchAnyPoda) {
            app.__bonsoSwitchAnyPoda = 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] no se pudo podar la rama inactiva",
                        e
                    );
                }
                return r;
            };
        }

        // --- Trabajo 2: "on true"/"on false" solo admiten el mismo tipo ---
        // Se busca por NOMBRE, no por indice: si el usuario convierte la
        // casilla `switch` a entrada (enlaza un booleano), ComfyUI la mete
        // en inputs y todo se desplaza. `switch` queda fuera de la
        // restriccion a proposito.
        const originalConnect = nodeType.prototype.onConnectInput;
        nodeType.prototype.onConnectInput = function (
            inputIndex,
            outputType,
            outputSlot,
            outputNode,
            outputSlotIndex
        ) {
            const sock = (this.inputs || [])[inputIndex];
            const nombre = sock ? sock.name : null;
            if (nombre !== "on true" && nombre !== "on false") {
                if (originalConnect) return originalConnect.apply(this, arguments);
                return true;
            }
            const otro = nombre === "on true" ? "on false" : "on true";
            const tipoOtro = tipoDeEntradaPorNombre(this, otro);
            if (
                tipoOtro &&
                tipoOtro !== "*" &&
                outputType &&
                outputType !== "*" &&
                tipoOtro !== outputType
            ) {
                console.debug(
                    `[Bonso Switch] tipo rechazado: '${outputType}' no coincide ` +
                        `con '${tipoOtro}' de la otra entrada`
                );
                return false;
            }
            if (originalConnect) return originalConnect.apply(this, arguments);
            return true;
        };
    },
});

function tipoDeEntradaPorNombre(node, nombre) {
    const idx = (node.inputs || []).findIndex((s) => s.name === nombre);
    if (idx < 0) return null;
    const sock = node.inputs[idx];
    if (!sock || sock.link == null) return null;
    const graph = node.graph;
    if (!graph || !graph.links) return null;
    const link =
        typeof graph.links.get === "function"
            ? graph.links.get(sock.link)
            : graph.links[sock.link];
    return link ? link.type : null;
}

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") continue;
        const inputs = nodo.inputs || (nodo.inputs = {});
        const activo = valorDeSwitch(output, inputs.switch);
        if (activo === null) continue;
        const rama = activo ? "on false" : "on true";
        if (rama in inputs) {
            delete inputs[rama];
            quitadas++;
        }
    }
    if (quitadas) {
        console.debug(
            `[Bonso Switch] rama inactiva fuera del prompt (${quitadas})`
        );
    }
    return quitadas;
}