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

const NODE_TYPE = "O1keyOmniFlashVideo";
const REFERENCE_GROUP = "参考图片";
const MEDIA_INPUTS = ["首帧图片", "尾帧图片", "源视频"];
const MODE_INPUTS = {
    "文生视频": [],
    "参考图视频": [REFERENCE_GROUP],
    "首尾帧": ["首帧图片", "尾帧图片"],
    "视频编辑": [REFERENCE_GROUP, "源视频"],
};

function isReferenceInput(input) {
    return input.name.startsWith(`${REFERENCE_GROUP}.`)
        || /^参考图片\d+$/.test(input.name);
}

function inputOptions(input) {
    return { label: input.label ?? input.name.split(".").at(-1), shape: input.shape };
}

function syncModeInputs(node) {
    const state = node._o1keyOmniModeInputs;
    if (!state || state.syncing) return;
    const visible = MODE_INPUTS[state.mode.value] ?? MODE_INPUTS["文生视频"];
    state.syncing = true;
    try {
        // Disable Autogrow before removing reference sockets; disconnects must
        // not recreate inputs belonging to the previous mode.
        if (visible.includes(REFERENCE_GROUP)) {
            node.comfyDynamic.autogrow[REFERENCE_GROUP] = state.autogrow;
        } else {
            delete node.comfyDynamic.autogrow[REFERENCE_GROUP];
        }
        for (let index = node.inputs.length - 1; index >= 0; index--) {
            const input = node.inputs[index];
            const active = isReferenceInput(input)
                ? visible.includes(REFERENCE_GROUP)
                : !MEDIA_INPUTS.includes(input.name) || visible.includes(input.name);
            if (!active) node.removeInput(index);
        }
        for (const name of visible) {
            const template = state.templates[name];
            const exists = name === REFERENCE_GROUP
                ? node.inputs.some(isReferenceInput)
                : node.inputs.some((input) => input.name === name);
            if (!exists) node.addInput(template.name, template.type, inputOptions(template));
        }
    } finally {
        state.syncing = false;
    }
    if (typeof node.computeSize === "function" && typeof node.setSize === "function") {
        const size = node.computeSize([...node.size]);
        node.setSize([node.size[0], size[1]]);
    }
    node.setDirtyCanvas?.(true, true);
}

function installModeInputs(node) {
    if ((node?.comfyClass || node?.type) !== NODE_TYPE || node._o1keyOmniModeInputs) return;
    const mode = node.widgets?.find((widget) => widget.name === "生成模式");
    const autogrow = node.comfyDynamic?.autogrow?.[REFERENCE_GROUP];
    if (!mode || !autogrow || !node.inputs || !node.addInput || !node.removeInput) return;
    const reference = node.inputs.find(isReferenceInput) ?? {
        name: `${REFERENCE_GROUP}.${autogrow.names?.[0] ?? "参考图片1"}`,
        type: "IMAGE",
    };
    const templates = { [REFERENCE_GROUP]: { ...reference } };
    for (const name of MEDIA_INPUTS) {
        const type = name === "源视频" ? "VIDEO" : "IMAGE";
        templates[name] = { ...(node.inputs.find((input) => input.name === name) ?? { name, type }) };
    }
    node._o1keyOmniModeInputs = { mode, autogrow, templates, syncing: false };
    const originalConnectionsChange = node.onConnectionsChange;
    node.onConnectionsChange = function () {
        if (node._o1keyOmniModeInputs.syncing) return;
        return originalConnectionsChange?.apply(this, arguments);
    };
    const originalCallback = mode.callback;
    mode.callback = function () {
        const result = originalCallback?.apply(this, arguments);
        syncModeInputs(node);
        return result;
    };
    syncModeInputs(node);
}

function installRunButton(node) {
    if ((node?.comfyClass || node?.type) !== NODE_TYPE || node._o1keyOmniRunButton) return;
    const button = node.addWidget("button", "开始生成", null, async () => {
        if (node._o1keyOmniSubmitting) return;
        node._o1keyOmniSubmitting = true;
        button.disabled = true;
        try {
            const queued = await app.queuePrompt(0, 1, [node.id]);
            if (queued === false) throw new Error("ComfyUI 未接受本次任务");
        } catch (error) {
            app.extensionManager?.toast?.add?.({
                severity: "error",
                summary: "Omni Flash 提交失败",
                detail: String(error?.message || error),
            });
        } finally {
            button.disabled = false;
            node._o1keyOmniSubmitting = false;
        }
    }, { serialize: false });
    button.serializeValue = () => undefined;
    node._o1keyOmniRunButton = button;
}

app.registerExtension({
    name: "o1key.omniFlashVideo",
    nodeCreated(node) {
        installRunButton(node);
        installModeInputs(node);
    },
    loadedGraphNode(node) {
        installRunButton(node);
        installModeInputs(node);
        syncModeInputs(node);
    },
});