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

const NODE_TYPES = new Set([
    "NanoBanana",
    "BatchNanoBananaPro",
    "O1keyGPTImage",
    "O1keyGPTImageBatch",
]);
const ROUTE_LABELS = Object.freeze({
    "畅速": "特价",
    "直连": "优质",
    "专线": "企业",
});
const INSTALLED_FLAG = "o1keyNanoBananaRouteLabelsInstalled";

function installRouteLabels(node) {
    if (!NODE_TYPES.has(node?.comfyClass)) return;

    const widget = node.widgets?.find(item => item.name === "模型线路");
    if (!widget?.options || widget.options[INSTALLED_FLAG]) return;

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

app.registerExtension({
    name: "o1key.nanoBananaRouteLabels",
    nodeCreated: installRouteLabels,
    loadedGraphNode: installRouteLabels,
});