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

const NODE_TYPES = new Set(["O1keyGPTImage", "O1keyGPTImageBatch"]);
const BASE_QUALITY_OPTIONS = ["高", "中", "低", "自动"];
const GPT_IMAGE_25_QUALITY_OPTIONS = [...BASE_QUALITY_OPTIONS, "超高", "最高"];
const QUALITY_FALLBACK = "自动";

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

function updateQualityOptions(node) {
    const modelWidget = findWidget(node, "模型");
    const qualityWidget = findWidget(node, "质量");
    if (!modelWidget || !qualityWidget) return;

    const supports25Quality =
        typeof modelWidget.value === "string" && modelWidget.value.startsWith("gpt-image-2.5-");
    const allowedOptions = supports25Quality ? GPT_IMAGE_25_QUALITY_OPTIONS : BASE_QUALITY_OPTIONS;
    qualityWidget.options ??= {};
    qualityWidget.options.values = [...allowedOptions];

    if (!allowedOptions.includes(qualityWidget.value)) {
        qualityWidget.value = QUALITY_FALLBACK;
        qualityWidget.callback?.(QUALITY_FALLBACK);
    }

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

function guardNode(node) {
    if (node.__o1keyGptImageQuality) return;
    const modelWidget = findWidget(node, "模型");
    const qualityWidget = findWidget(node, "质量");
    if (!modelWidget || !qualityWidget) return;
    node.__o1keyGptImageQuality = true;

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

    updateQualityOptions(node);
}

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

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

    loadedGraphNode(node) {
        if (!NODE_TYPES.has(node.comfyClass || node.type)) return;
        guardNode(node);
        updateQualityOptions(node);
    },
});