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

const NODE_TYPES = new Set(["BatchNanoBananaPro", "O1keyGPTImageBatch"]);
const PATH_COUNT_WIDGET_NAMES = new Set(["图片路径数量", "图片文件夹数量"]);
const AUTOGROW_INPUT_PATTERN = /^参考图组\.参考图(\d+)$/;

function parsePathCount(value) {
    const count = Number.parseInt(String(value ?? ""), 10);
    return Number.isFinite(count) ? Math.max(1, Math.min(5, count)) : 1;
}

function findPathCountWidget(node) {
    return node.widgets?.find((widget) => PATH_COUNT_WIDGET_NAMES.has(widget.name));
}

function updateReferenceInputLabels(node) {
    const widget = findPathCountWidget(node);
    const pathCount = parsePathCount(widget?.value);
    let changed = false;

    for (const input of node.inputs ?? []) {
        const match = AUTOGROW_INPUT_PATTERN.exec(input?.name ?? "");
        if (!match) continue;
        const label = `参考图${pathCount + Number(match[1])}`;
        if (input.label !== label) {
            // 只改画布显示标签，不改稳定的内部端口名，避免工作流连线失效。
            input.label = label;
            changed = true;
        }
    }

    if (changed) node.setDirtyCanvas?.(true, true);
}

function scheduleReferenceInputLabels(node) {
    updateReferenceInputLabels(node);
    requestAnimationFrame(() => updateReferenceInputLabels(node));
}

function guardNode(node) {
    if (node.__o1keyBatchReferenceImageLabels || node.__o1keyBatchNanoBananaReferenceLabels) return;
    node.__o1keyBatchReferenceImageLabels = true;

    const bindPathCountWidget = () => {
        const widget = findPathCountWidget(node);
        if (!widget || widget.__o1keyReferenceLabelCallback) return;
        widget.__o1keyReferenceLabelCallback = true;
        const originalCallback = widget.callback;
        widget.callback = function () {
            const result = originalCallback?.apply(this, arguments);
            scheduleReferenceInputLabels(node);
            return result;
        };
    };

    const originalOnConnectionsChange = node.onConnectionsChange;
    node.onConnectionsChange = function () {
        const result = originalOnConnectionsChange?.apply(this, arguments);
        scheduleReferenceInputLabels(this);
        return result;
    };

    const originalOnConfigure = node.onConfigure;
    node.onConfigure = function () {
        const result = originalOnConfigure?.apply(this, arguments);
        bindPathCountWidget();
        scheduleReferenceInputLabels(this);
        return result;
    };

    bindPathCountWidget();
    scheduleReferenceInputLabels(node);
}

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

    nodeCreated(node) {
        if (NODE_TYPES.has(node.comfyClass)) guardNode(node);
    },

    loadedGraphNode(node) {
        if (NODE_TYPES.has(node.comfyClass)) {
            guardNode(node);
            scheduleReferenceInputLabels(node);
        }
    },
});