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

const NODE_TYPES = new Set(["O1keyGPTImage", "O1keyGPTImageBatch"]);
const BACKGROUND_LABELS = Object.freeze({
    auto: "自动",
    transparent: "透明",
    opaque: "不透明",
});
const INSTALLED_FLAG = "o1keyGptImageBackgroundLabelsInstalled";

function installBackgroundLabels(node) {
    if (!NODE_TYPES.has(node?.comfyClass || node?.type)) return;

    const widget = node.widgets?.find((item) => item.name === "背景");
    if (!widget) return;
    widget.options ??= {};
    if (widget.options[INSTALLED_FLAG]) return;

    const previousGetOptionLabel = widget.options.getOptionLabel;
    widget.options.getOptionLabel = (value) => {
        const normalized = value == null ? "" : String(value);
        return BACKGROUND_LABELS[normalized]
            ?? previousGetOptionLabel?.(value)
            ?? normalized;
    };
    widget.options[INSTALLED_FLAG] = true;
    node.setDirtyCanvas?.(true, false);
}

app.registerExtension({
    name: "o1key.gptImageBackgroundLabels",
    nodeCreated: installBackgroundLabels,
    loadedGraphNode: installBackgroundLabels,
});