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

const NODE_TYPE = "O1keyPromptMultiFunction";
const MODE_WIDGET = "功能";
const SAMPLE_COUNT_WIDGET = "抽取数量";
const SELECTED_INDICES_WIDGET = "指定序号";
const RANDOM_MODES = new Set(["随机抽取n套", "随机抽取1套", "随机抽取多套"]);
const SELECTED_MODE = "指定序号";

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

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

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

function syncPromptMultiFunctionWidgets(node) {
    const modeWidget = findWidget(node, MODE_WIDGET);
    const sampleCountWidget = findWidget(node, SAMPLE_COUNT_WIDGET);
    const selectedIndicesWidget = findWidget(node, SELECTED_INDICES_WIDGET);
    if (!modeWidget || !sampleCountWidget || !selectedIndicesWidget) return;

    const randomMode = RANDOM_MODES.has(modeWidget.value);
    const selectedMode = modeWidget.value === SELECTED_MODE;
    const countChanged = setWidgetHidden(sampleCountWidget, !randomMode);
    const indicesChanged = setWidgetHidden(selectedIndicesWidget, !selectedMode);
    if (countChanged || indicesChanged) node.setDirtyCanvas?.(true, true);
}

function scheduleWidgetSync(node) {
    syncPromptMultiFunctionWidgets(node);
    requestAnimationFrame(() => syncPromptMultiFunctionWidgets(node));
}

function guardNode(node) {
    if (node.__o1keyPromptMultiFunctionDynamic) return;
    if (!findWidget(node, MODE_WIDGET)
        || !findWidget(node, SAMPLE_COUNT_WIDGET)
        || !findWidget(node, SELECTED_INDICES_WIDGET)) return;
    node.__o1keyPromptMultiFunctionDynamic = true;

    const bindModeWidget = () => {
        const modeWidget = findWidget(node, MODE_WIDGET);
        if (!modeWidget || modeWidget.__o1keyPromptMultiFunctionCallback) return;
        modeWidget.__o1keyPromptMultiFunctionCallback = true;
        const originalCallback = modeWidget.callback;
        modeWidget.callback = function () {
            const result = originalCallback?.apply(this, arguments);
            scheduleWidgetSync(node);
            return result;
        };
    };

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

    bindModeWidget();
    scheduleWidgetSync(node);
}

function isTargetNode(node) {
    return node?.comfyClass === NODE_TYPE || node?.type === NODE_TYPE;
}

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

    nodeCreated(node) {
        if (isTargetNode(node)) guardNode(node);
    },

    loadedGraphNode(node) {
        if (!isTargetNode(node)) return;
        guardNode(node);
        scheduleWidgetSync(node);
    },
});