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

const NODE_TYPES = new Set(["SeedanceAutoPass", "SeedanceAutoPassBatch"]);
const MODEL_25 = "seedance 2.5";
const OVERSEAS_ROUTE = "海外";
const LEGACY_MODEL_ROUTES = new Set(["海外HC", "海外破限高并发", "海外破限", "海外破限标准", "海外标准"]);
const ALL_RESOLUTIONS = ["480p", "720p", "1080p", "4k"];
const LIMITED_RESOLUTIONS = ["480p", "720p"];
const LIMITED_RESOLUTION_MODELS = new Set([
    "seedance 2.0 fast",
    "seedance 2.0 mini",
]);
const RESOLUTION_FALLBACK = "720p";
const ID_WIDGET_GROUPS = [
    { prefix: "图片素材ID", maximum: 30 },
    { prefix: "视频素材ID", maximum: 10 },
    { prefix: "音频素材ID", maximum: 10 },
];
const UNAVAILABLE_WIDGETS = ["联网搜索", "返回末帧图片"];

const durationOptions = (maximum) => [
    "自动",
    ...Array.from({ length: maximum - 3 }, (_, index) => `${index + 4}秒`),
];

const DURATIONS_20 = durationOptions(15);
const DURATIONS_25 = durationOptions(30);

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

function findAssetModeWidget(node) {
    return findWidget(node, "素材创建模式") ?? findWidget(node, "素材创建");
}

const REFERENCE_GROUPS = ["参考图片", "参考视频", "参考音频"];
const FRAME_INPUTS = ["首帧图片", "尾帧图片"];

function isReferenceInput(input) {
    return REFERENCE_GROUPS.some((group) => input.name.startsWith(`${group}.`)
        || new RegExp(`^${group}\\d+$`).test(input.name));
}

function inputOptions(input) {
    // Never reuse a removed socket's link or internal identity.
    return { label: input.label ?? input.name.split(".").at(-1), shape: input.shape };
}

function syncModeInputs(node) {
    const state = node.__o1keyAutoPassModeInputs;
    if (!state || state.syncing) return;
    const frameMode = findWidget(node, "生成模式")?.value === "首尾帧";
    const autogrow = node.comfyDynamic.autogrow;
    state.syncing = true;
    try {
        // Suspend native Autogrow before removing its sockets; otherwise a
        // disconnect callback can recreate an inactive reference input.
        for (const group of REFERENCE_GROUPS) {
            if (frameMode) delete autogrow[group];
            else autogrow[group] = state.groups[group];
        }
        for (let index = node.inputs.length - 1; index >= 0; index--) {
            const input = node.inputs[index];
            if (frameMode ? isReferenceInput(input) : FRAME_INPUTS.includes(input.name)) {
                node.removeInput(index);
            }
        }
        const templates = frameMode ? state.frames : state.references;
        for (const template of templates) {
            const exists = frameMode
                ? node.inputs.some((input) => input.name === template.name)
                : node.inputs.some((input) => input.name.startsWith(`${template.name.split(".")[0]}.`));
            if (!exists) node.addInput(template.name, template.type, inputOptions(template));
        }
    } finally {
        state.syncing = false;
    }
    syncIdWidgets(node);
}

function guardModeInputs(node) {
    if (node.comfyClass !== "SeedanceAutoPass" || node.__o1keyAutoPassModeInputs) return;
    const mode = findWidget(node, "生成模式");
    const autogrow = node.comfyDynamic?.autogrow;
    if (!mode || !autogrow || !node.addInput || !node.removeInput) return;
    const references = REFERENCE_GROUPS.map((group) => node.inputs?.find((input) => input.name.startsWith(`${group}.`)));
    if (references.some((input) => !input) || REFERENCE_GROUPS.some((group) => !autogrow[group])) return;
    node.__o1keyAutoPassModeInputs = {
        syncing: false,
        groups: Object.fromEntries(REFERENCE_GROUPS.map((group) => [group, autogrow[group]])),
        references: references.map((input) => ({ ...input })),
        frames: FRAME_INPUTS.map((name) => ({ ...(node.inputs.find((input) => input.name === name) ?? { name, type: "IMAGE" }) })),
    };
    const originalConnectionsChange = node.onConnectionsChange;
    node.onConnectionsChange = function () {
        if (node.__o1keyAutoPassModeInputs.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);
}

// 与 SeedanceMultiModal 相同：保留控件和值的序列化顺序，仅渐进显隐单行 ID。
function setIdWidgetHidden(widget, hidden) {
    if (!widget || (widget.__o1keyAutoPassIdHidden === hidden
        && widget.hidden === hidden && widget.options?.hidden === hidden)) return;
    widget.__o1keyAutoPassIdHidden = hidden;
    widget.hidden = hidden;
    widget.options ??= {};
    widget.options.hidden = hidden;
    if (!("__o1keyAutoPassOriginalComputeSize" in widget)) {
        widget.__o1keyAutoPassOriginalComputeSize = widget.computeSize ?? null;
    }
    if (hidden) {
        widget.computeSize = () => [0, -4];
    } else if (widget.__o1keyAutoPassOriginalComputeSize) {
        widget.computeSize = widget.__o1keyAutoPassOriginalComputeSize;
    } else {
        delete widget.computeSize;
    }
}

function syncIdWidgets(node) {
    if (node.comfyClass !== "SeedanceAutoPass") return;
    // Temporarily unavailable controls stay serialized for saved workflows.
    for (const name of UNAVAILABLE_WIDGETS) setIdWidgetHidden(findWidget(node, name), true);
    const assetMode = findAssetModeWidget(node);
    // 该开关自身始终显示，只控制下方 ID 行。
    if (assetMode) {
        assetMode.hidden = false;
        assetMode.options ??= {};
        assetMode.options.hidden = false;
    }
    const manual = ["打开", "手动", "manual"].includes(assetMode?.value);
    for (const { prefix, maximum } of ID_WIDGET_GROUPS) {
        const widgets = Array.from({ length: maximum }, (_, index) => findWidget(node, `${prefix}${index + 1}`));
        let highestFilled = -1;
        widgets.forEach((widget, index) => {
            if (typeof widget?.value === "string" && widget.value.trim()) highestFilled = index;
        });
        const visibleCount = manual ? Math.min(maximum, Math.max(1, highestFilled + 2)) : 0;
        widgets.forEach((widget, index) => setIdWidgetHidden(widget, index >= visibleCount));
    }
    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 guardIdWidgets(node) {
    if (node.comfyClass !== "SeedanceAutoPass" || node.__o1keyAutoPassIdGuard) return;
    node.__o1keyAutoPassIdGuard = true;
    const widgets = [findAssetModeWidget(node)];
    for (const { prefix, maximum } of ID_WIDGET_GROUPS) {
        for (let index = 1; index <= maximum; index++) widgets.push(findWidget(node, `${prefix}${index}`));
    }
    for (const widget of widgets) {
        if (!widget) continue;
        const originalCallback = widget.callback;
        widget.callback = function () {
            const result = originalCallback?.apply(this, arguments);
            syncIdWidgets(node);
            return result;
        };
    }
    syncIdWidgets(node);
}

function notify(detail) {
    app.extensionManager?.toast?.add({
        severity: "info",
        summary: "Seedance 全能生成视频参数已调整",
        detail,
        life: 3500,
    });
}

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

function updateDynamicParameters(node, silent = false) {
    const modelWidget = findWidget(node, "主模型");
    const routeWidget = findWidget(node, "模型线路");
    const resolutionWidget = findWidget(node, "分辨率");
    const durationWidget = findWidget(node, "时长");
    if (!modelWidget || !routeWidget || !resolutionWidget || !durationWidget) return;

    const is25 = modelWidget.value === MODEL_25;
    const hasLimitedResolution = LIMITED_RESOLUTION_MODELS.has(modelWidget.value);
    const maximumDuration = is25 ? 30 : 15;
    durationWidget.options ??= {};
    durationWidget.options.values = is25 ? DURATIONS_25 : DURATIONS_20;
    resolutionWidget.options ??= {};
    resolutionWidget.options.values = hasLimitedResolution ? LIMITED_RESOLUTIONS : ALL_RESOLUTIONS;

    if (LEGACY_MODEL_ROUTES.has(routeWidget.value)) {
        setWidgetValue(routeWidget, OVERSEAS_ROUTE);
    }

    if (hasLimitedResolution && !LIMITED_RESOLUTIONS.includes(resolutionWidget.value)) {
        setWidgetValue(resolutionWidget, RESOLUTION_FALLBACK);
        if (!silent) {
            notify(`${modelWidget.value} 仅支持 480p/720p，已自动切换为 ${RESOLUTION_FALLBACK}。`);
        }
    }

    const duration = Number.parseInt(durationWidget.value, 10);
    if (Number.isFinite(duration) && duration > maximumDuration) {
        setWidgetValue(durationWidget, `${maximumDuration}秒`);
        if (!silent) {
            notify(`${modelWidget.value} 的最大时长为 ${maximumDuration} 秒，已自动调整。`);
        }
    }

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

function guardNode(node) {
    if (node.__o1keySeedanceAutoPassDynamic) return;
    node.__o1keySeedanceAutoPassDynamic = true;

    const modelWidget = findWidget(node, "主模型");
    const routeWidget = findWidget(node, "模型线路");
    const resolutionWidget = findWidget(node, "分辨率");
    const durationWidget = findWidget(node, "时长");
    if (!modelWidget || !routeWidget || !resolutionWidget || !durationWidget) return;

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

    const originalRouteCallback = routeWidget.callback;
    routeWidget.callback = function (value) {
        if (LEGACY_MODEL_ROUTES.has(value)) {
            routeWidget.value = OVERSEAS_ROUTE;
            node.setDirtyCanvas?.(true, true);
            return originalRouteCallback?.call(this, OVERSEAS_ROUTE);
        }
        return originalRouteCallback?.apply(this, arguments);
    };

    const originalResolutionCallback = resolutionWidget.callback;
    resolutionWidget.callback = function (value) {
        if (LIMITED_RESOLUTION_MODELS.has(modelWidget.value) && !LIMITED_RESOLUTIONS.includes(value)) {
            resolutionWidget.value = RESOLUTION_FALLBACK;
            notify(`${modelWidget.value} 仅支持 480p/720p。`);
            node.setDirtyCanvas?.(true, true);
            return originalResolutionCallback?.call(this, RESOLUTION_FALLBACK);
        }
        return originalResolutionCallback?.apply(this, arguments);
    };

    const originalDurationCallback = durationWidget.callback;
    durationWidget.callback = function (value) {
        const maximumDuration = modelWidget.value === MODEL_25 ? 30 : 15;
        const duration = Number.parseInt(value, 10);
        if (Number.isFinite(duration) && duration > maximumDuration) {
            const fallback = `${maximumDuration}秒`;
            durationWidget.value = fallback;
            notify(`${modelWidget.value} 的最大时长为 ${maximumDuration} 秒。`);
            node.setDirtyCanvas?.(true, true);
            return originalDurationCallback?.call(this, fallback);
        }
        return originalDurationCallback?.apply(this, arguments);
    };

    updateDynamicParameters(node, true);
}

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

    nodeCreated(node) {
        if (!NODE_TYPES.has(node.comfyClass)) return;
        guardNode(node);
        guardIdWidgets(node);
        guardModeInputs(node);
    },

    loadedGraphNode(node) {
        if (!NODE_TYPES.has(node.comfyClass)) return;
        updateDynamicParameters(node, true);
        guardIdWidgets(node);
        guardModeInputs(node);
        syncModeInputs(node);
        syncIdWidgets(node);
    },
});