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

const NODE_TYPES = new Set(["NanoBanana", "BatchNanoBananaPro"]);
const MODEL_WIDGET = "模型";
const THINKING_WIDGET = "思考等级";
const SUPPORTED_MODEL = "Nano Banana 2";

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

function setThinkingLevelVisibility(node) {
    const modelWidget = findWidget(node, MODEL_WIDGET);
    const thinkingWidget = findWidget(node, THINKING_WIDGET);
    if (!modelWidget || !thinkingWidget) return;

    const hidden = modelWidget.value !== SUPPORTED_MODEL;
    if (thinkingWidget.__o1keyThinkingHidden === hidden) return;
    thinkingWidget.__o1keyThinkingHidden = hidden;

    // Vue 节点使用 options.hidden；旧版 LiteGraph 使用 hidden/computeSize。
    // 两种状态同时更新，确保隐藏后既不绘制，也不留下空白行。
    thinkingWidget.hidden = hidden;
    thinkingWidget.options ??= {};
    thinkingWidget.options.hidden = hidden;

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

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

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

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

    const bindModelWidget = () => {
        const modelWidget = findWidget(node, MODEL_WIDGET);
        if (!modelWidget || modelWidget.__o1keyThinkingLevelCallback) return;

        modelWidget.__o1keyThinkingLevelCallback = true;
        const originalCallback = modelWidget.callback;
        modelWidget.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);
        bindModelWidget();
        scheduleVisibilityUpdate(this);
        return result;
    };

    bindModelWidget();
    scheduleVisibilityUpdate(node);
}

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

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

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