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

const EXTENSION = "ethanfel.prompt_builder.index_switch_slots";
const NODE_NAME = "SxCPIndexSwitch";
const MAX_INPUTS = 64;

function isSwitchInput(input) {
    return /^input_\d+$/.test(input?.name || "");
}

function isSwitchOutput(output) {
    return /^output_\d+$/.test(output?.name || "");
}

function slotNumber(slot) {
    const match = String(slot?.name || "").match(/\d+$/);
    return match ? Number(match[0]) : -1;
}

function resizeNode(node) {
    const size = node.computeSize?.();
    if (size) node.setSize?.(size);
    app.graph?.setDirtyCanvas(true, true);
}

function addSwitchInput(node, number) {
    if (number < 1 || number > MAX_INPUTS) return;
    const name = `input_${number}`;
    if (!node.inputs?.some((input) => input.name === name)) {
        node.addInput(name, "*");
    }
}

function removeSwitchInput(node, number) {
    const inputIndex = node.inputs?.findIndex((input) => input.name === `input_${number}`) ?? -1;
    if (inputIndex >= 0 && !node.inputs[inputIndex]?.link) {
        node.removeInput(inputIndex);
    }
}

function addSwitchOutput(node, number) {
    if (number < 1 || number > MAX_INPUTS) return;
    const name = `output_${number}`;
    if (!node.outputs?.some((output) => output.name === name)) {
        node.addOutput(name, "*");
    }
}

function removeSwitchOutput(node, number) {
    const outputIndex = node.outputs?.findIndex((output) => output.name === `output_${number}`) ?? -1;
    if (outputIndex >= 0 && !(node.outputs[outputIndex]?.links?.length)) {
        node.removeOutput(outputIndex);
    }
}

function trimInputTail(node) {
    for (let number = MAX_INPUTS; number > 1; number--) {
        const input = node.inputs?.find((slot) => slot.name === `input_${number}`);
        const previous = node.inputs?.find((slot) => slot.name === `input_${number - 1}`);
        if (!input?.link && !previous?.link) removeSwitchInput(node, number);
    }
}

function trimOutputTail(node) {
    for (let number = MAX_INPUTS; number > 1; number--) {
        const output = node.outputs?.find((slot) => slot.name === `output_${number}`);
        const previous = node.outputs?.find((slot) => slot.name === `output_${number - 1}`);
        if (!(output?.links?.length) && !(previous?.links?.length)) removeSwitchOutput(node, number);
    }
}

function setupNodeSlots(node) {
    addSwitchInput(node, 1);
    addSwitchOutput(node, 1);
    for (let number = 2; number <= MAX_INPUTS; number++) {
        const input = node.inputs?.find((slot) => slot.name === `input_${number}`);
        if (!input?.link) removeSwitchInput(node, number);
        const output = node.outputs?.find((slot) => slot.name === `output_${number}`);
        if (!(output?.links?.length)) removeSwitchOutput(node, number);
    }
    trimInputTail(node);
    trimOutputTail(node);
    resizeNode(node);
}

function maybeGrowInput(node) {
    const switchInputs = (node.inputs || []).filter(isSwitchInput);
    const last = switchInputs.reduce((max, input) => Math.max(max, slotNumber(input)), 1);
    const lastInput = node.inputs?.find((slot) => slot.name === `input_${last}`);
    if (lastInput?.link && last < MAX_INPUTS) {
        addSwitchInput(node, last + 1);
        resizeNode(node);
    }
}

function maybeGrowOutput(node) {
    const switchOutputs = (node.outputs || []).filter(isSwitchOutput);
    const last = switchOutputs.reduce((max, output) => Math.max(max, slotNumber(output)), 1);
    const lastOutput = node.outputs?.find((slot) => slot.name === `output_${last}`);
    if (lastOutput?.links?.length && last < MAX_INPUTS) {
        addSwitchOutput(node, last + 1);
        resizeNode(node);
    }
}

app.registerExtension({
    name: EXTENSION,

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

        const onNodeCreated = nodeType.prototype.onNodeCreated;
        nodeType.prototype.onNodeCreated = function () {
            const result = onNodeCreated?.apply(this, arguments);
            queueMicrotask(() => setupNodeSlots(this));
            return result;
        };

        const onConfigure = nodeType.prototype.onConfigure;
        nodeType.prototype.onConfigure = function () {
            const result = onConfigure?.apply(this, arguments);
            queueMicrotask(() => setupNodeSlots(this));
            return result;
        };

        const onConnectionsChange = nodeType.prototype.onConnectionsChange;
        nodeType.prototype.onConnectionsChange = function (type, index, connected, linkInfo) {
            const result = onConnectionsChange?.apply(this, arguments);
            if (!linkInfo) return result;
            const slot = type === LiteGraph.INPUT ? this.inputs?.[index] : this.outputs?.[index];
            if (type === LiteGraph.INPUT && isSwitchInput(slot)) {
                if (connected) maybeGrowInput(this);
                else {
                    trimInputTail(this);
                    resizeNode(this);
                }
            } else if (isSwitchOutput(slot)) {
                if (connected) maybeGrowOutput(this);
                else {
                    trimOutputTail(this);
                    resizeNode(this);
                }
            }
            return result;
        };
    },
});