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

const NODE_TYPE = "SeedanceMultiModal";
const ID_WIDGET_GROUPS = [
    { prefix: "图片素材ID", maximum: 30 },
    { prefix: "视频素材ID", maximum: 10 },
    { prefix: "音频素材ID", maximum: 10 },
];

function findWidget(node, name) {
    return node.widgets?.find((widget) => widget.name === name);
}

function hasValue(widget) {
    return typeof widget?.value === "string"
        ? widget.value.trim().length > 0
        : widget?.value != null;
}

function setWidgetHidden(widget, hidden) {
    if (!widget || widget.__o1keyProgressiveHidden === hidden) return;
    widget.__o1keyProgressiveHidden = hidden;
    widget.hidden = hidden;
    widget.options ??= {};
    widget.options.hidden = hidden;

    if (!("__o1keyOriginalComputeSize" in widget)) {
        widget.__o1keyOriginalComputeSize = widget.computeSize ?? null;
    }
    if (hidden) {
        widget.computeSize = () => [0, -4];
    } else if (widget.__o1keyOriginalComputeSize) {
        widget.computeSize = widget.__o1keyOriginalComputeSize;
    } else {
        delete widget.computeSize;
    }
}

function resizeNodeToVisibleWidgets(node) {
    if (typeof node.computeSize !== "function" || typeof node.setSize !== "function") return;
    const computedSize = node.computeSize([...node.size]);
    node.setSize([node.size[0], computedSize[1]]);
}

function syncProgressiveIdWidgets(node) {
    for (const { prefix, maximum } of ID_WIDGET_GROUPS) {
        const widgets = Array.from(
            { length: maximum },
            (_, index) => findWidget(node, `${prefix}${index + 1}`),
        );
        let highestFilled = -1;
        widgets.forEach((widget, index) => {
            if (hasValue(widget)) highestFilled = index;
        });
        const visibleCount = Math.min(maximum, Math.max(1, highestFilled + 2));
        widgets.forEach((widget, index) => setWidgetHidden(widget, index >= visibleCount));
    }
    resizeNodeToVisibleWidgets(node);
    node.setDirtyCanvas?.(true, true);
}

function guardNode(node) {
    if (node.__o1keySeedanceMultiModalDynamic) return;
    node.__o1keySeedanceMultiModalDynamic = true;

    for (const { prefix, maximum } of ID_WIDGET_GROUPS) {
        for (let index = 1; index <= maximum; index++) {
            const widget = findWidget(node, `${prefix}${index}`);
            if (!widget || widget.__o1keyProgressiveCallback) continue;
            widget.__o1keyProgressiveCallback = true;
            const originalCallback = widget.callback;
            widget.callback = function () {
                const result = originalCallback?.apply(this, arguments);
                syncProgressiveIdWidgets(node);
                return result;
            };
        }
    }

    syncProgressiveIdWidgets(node);
}

app.registerExtension({
    name: "o1key.seedanceMultiModalDynamic",

    nodeCreated(node) {
        if (node.comfyClass === NODE_TYPE) guardNode(node);
    },

    loadedGraphNode(node) {
        if (node.comfyClass !== NODE_TYPE) return;
        guardNode(node);
        syncProgressiveIdWidgets(node);
    },
});