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

const NODE_TYPE = "BatchNanoBananaPro";
const FORMAT_WIDGET = "图片输出格式";
const QUALITY_WIDGET = "图片质量";
const LOSSY_FORMATS = new Set(["JPEG", "WebP"]);

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

function setQualityVisibility(node) {
    const formatWidget = findWidget(node, FORMAT_WIDGET);
    const qualityWidget = findWidget(node, QUALITY_WIDGET);
    if (!formatWidget || !qualityWidget) return;

    const numericQuality = Number(qualityWidget.value);
    qualityWidget.value = Number.isInteger(numericQuality)
        && numericQuality >= 1
        && numericQuality <= 100
        ? numericQuality
        : 95;

    const hidden = !LOSSY_FORMATS.has(formatWidget.value);
    if (qualityWidget.__o1keyImageQualityHidden === hidden) return;
    qualityWidget.__o1keyImageQualityHidden = hidden;

    qualityWidget.hidden = hidden;
    qualityWidget.options ??= {};
    qualityWidget.options.hidden = hidden;

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

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

function scheduleVisibilityUpdate(node) {
    setQualityVisibility(node);
    requestAnimationFrame(() => setQualityVisibility(node));
}

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

    const bindFormatWidget = () => {
        const formatWidget = findWidget(node, FORMAT_WIDGET);
        if (!formatWidget || formatWidget.__o1keyImageQualityCallback) return;

        formatWidget.__o1keyImageQualityCallback = true;
        const originalCallback = formatWidget.callback;
        formatWidget.callback = function () {
            const result = originalCallback?.apply(this, arguments);
            scheduleVisibilityUpdate(node);
            return result;
        };
    };

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

    bindFormatWidget();
    scheduleVisibilityUpdate(node);
}

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

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

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