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

const NODE_TYPE = "MiniMaxH3Video";
const MODEL_H3 = "MiniMax-H3";
const MODEL_MAX = "MiniMax-H3-MAX";
const MODE_TEXT = "文生视频";
const MODE_REFERENCE = "参考素材生视频";
const RESOLUTIONS = {
    [MODEL_H3]: ["2K", "768P"],
    [MODEL_MAX]: ["768P", "480P"],
};
const MIN_DURATION = {
    [MODEL_H3]: 4,
    [MODEL_MAX]: 5,
};

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

function notify(detail) {
    app.extensionManager?.toast?.add({
        severity: "warn",
        summary: "MiniMax H3 参数已调整",
        detail,
        life: 4000,
    });
}

function setWidgetValue(widget, value) {
    if (!widget || widget.value === value) return;
    widget.value = value;
    widget.callback?.(value);
}

function updateParameters(node, silent = false) {
    const modelWidget = findWidget(node, "模型");
    const modeWidget = findWidget(node, "生成模式");
    const resolutionWidget = findWidget(node, "分辨率");
    const durationWidget = findWidget(node, "时长");
    if (!modelWidget || !modeWidget || !resolutionWidget || !durationWidget) return;

    const model = modelWidget.value === MODEL_MAX ? MODEL_MAX : MODEL_H3;
    const allowedResolutions = RESOLUTIONS[model];
    resolutionWidget.options ??= {};
    resolutionWidget.options.values = [...allowedResolutions];
    durationWidget.options ??= {};
    durationWidget.options.min = MIN_DURATION[model];
    durationWidget.options.max = 15;

    if (!allowedResolutions.includes(resolutionWidget.value)) {
        setWidgetValue(resolutionWidget, "768P");
        if (!silent) notify(`${model} 不支持原分辨率，已切换为 768P。`);
    }

    const duration = Number(durationWidget.value);
    if (Number.isFinite(duration) && duration < MIN_DURATION[model]) {
        setWidgetValue(durationWidget, MIN_DURATION[model]);
        if (!silent) notify(`${model} 最短时长为 ${MIN_DURATION[model]} 秒。`);
    }

    if (model === MODEL_MAX && modeWidget.value === MODE_REFERENCE) {
        setWidgetValue(modeWidget, MODE_TEXT);
        if (!silent) notify("MiniMax-H3-MAX 不支持参考素材模式，已切换为文生视频。");
    }

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

function guardNode(node) {
    if (node.__o1keyMiniMaxH3ParameterGuard) return;

    const modelWidget = findWidget(node, "模型");
    const modeWidget = findWidget(node, "生成模式");
    const resolutionWidget = findWidget(node, "分辨率");
    const durationWidget = findWidget(node, "时长");
    if (!modelWidget || !modeWidget || !resolutionWidget || !durationWidget) return;
    node.__o1keyMiniMaxH3ParameterGuard = true;

    const originalModelCallback = modelWidget.callback;
    modelWidget.callback = function () {
        const result = originalModelCallback?.apply(this, arguments);
        updateParameters(node);
        return result;
    };

    const originalModeCallback = modeWidget.callback;
    modeWidget.callback = function (value) {
        if (modelWidget.value === MODEL_MAX && value === MODE_REFERENCE) {
            modeWidget.value = MODE_TEXT;
            notify("MiniMax-H3-MAX 不支持参考素材模式。");
            node.setDirtyCanvas?.(true, true);
            return originalModeCallback?.call(this, MODE_TEXT);
        }
        return originalModeCallback?.apply(this, arguments);
    };

    const originalResolutionCallback = resolutionWidget.callback;
    resolutionWidget.callback = function (value) {
        const model = modelWidget.value === MODEL_MAX ? MODEL_MAX : MODEL_H3;
        if (!RESOLUTIONS[model].includes(value)) {
            resolutionWidget.value = "768P";
            notify(`${model} 不支持 ${value}，已切换为 768P。`);
            node.setDirtyCanvas?.(true, true);
            return originalResolutionCallback?.call(this, "768P");
        }
        return originalResolutionCallback?.apply(this, arguments);
    };

    const originalDurationCallback = durationWidget.callback;
    durationWidget.callback = function (value) {
        const model = modelWidget.value === MODEL_MAX ? MODEL_MAX : MODEL_H3;
        const minimum = MIN_DURATION[model];
        const duration = Number(value);
        if (Number.isFinite(duration) && duration < minimum) {
            durationWidget.value = minimum;
            notify(`${model} 最短时长为 ${minimum} 秒。`);
            node.setDirtyCanvas?.(true, true);
            return originalDurationCallback?.call(this, minimum);
        }
        return originalDurationCallback?.apply(this, arguments);
    };

    updateParameters(node, true);
}

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

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

    loadedGraphNode(node) {
        if ((node.comfyClass || node.type) !== NODE_TYPE) return;
        guardNode(node);
        updateParameters(node, true);
    },
});