import { app } from "../../scripts/app.js";
import { collectChain, MODEL_COMBO_TYPES, STACK_TYPES, FLAKE_COMBO_TYPES } from "./widgets/generation-data.js";
import { serializeModelOverrides, expandEntryVariants } from "./utils.js";
import { fetchFlakeMeta, cachedAutoCaption, variantComboCount } from "./api.js";

// Full set of variant selections for a flake whose groups are `meta`
// ({group: [choices]}) — the cartesian product, one {group: choice} per combo.
// Empty meta yields a single empty selection (#343).
function enumerateVariantCombos(meta) {
    const groups = Object.keys(meta || {}).filter(g => Array.isArray(meta[g]) && meta[g].length);
    if (groups.length === 0) return [{}];
    const arrays = groups.map(g => meta[g].map(ch => [g, ch]));
    return cartesianProduct(arrays).map(pairs => Object.fromEntries(pairs));
}

// A FlakeStack's flakes_json entries.
function stackEntries(node) {
    const w = node.widgets?.find(w => w.name === "flakes_json");
    let entries = [];
    try { entries = JSON.parse(w?.value || "[]"); } catch { entries = []; }
    return Array.isArray(entries) ? entries : [];
}

// Active (non-bypassed, non-inline) "All Variants" flakes in a FlakeStack's
// flakes_json, as [{ index, flake }] (#343).
function stackAllVariantsFlakes(node) {
    return stackEntries(node)
        .map((flake, index) => ({ flake, index }))
        .filter(({ flake }) => flake && !flake.inline && !flake.bypassed && flake._all_variants && flake.name);
}

// Flake Animate grid items carry input images (#441), one job each. A combo
// item splits into [{ image, key: `${key}@${k}` }]; without images it stays
// one job under its own key. The Generation Data overlay mirrors these keys.
export function itemImageJobs(entry, key) {
    const images = entry.images || [];
    return images.length ? images.map((image, k) => ({ image, key: `${key}@${k}` })) : [{ image: null, key }];
}

// A Flake Animate Stack's jobs: one per input image of its active items, in
// grid order, as [{ index, image, key: `s@${k}` }] (index: the item's entry).
export function stackImageJobs(node) {
    return stackEntries(node)
        .flatMap((e, index) => (e && !e.bypassed && e.images || []).map(image => ({ index, image })))
        .map((job, k) => ({ ...job, key: `s@${k}` }));
}

export function getComboFlakes(node) {
    return node.properties?._combo_flakes || [];
}

export function getComboPresets(node) {
    return node.properties?._combo_presets || [];
}

// Persistent id of a combo grid item (#383): Flake Combo entries carry it as
// `_uid`, Flake Model Combo keeps `_combo_uids` parallel to `_combo_presets`.
// Combo keys use it, so results, previews and highlights survive reordering.
export function newItemUid() {
    return Math.random().toString(36).slice(2, 10);
}

// Sub-key of an "All Variants" item: its selection as sorted group=choice pairs.
export function variantKey(variant) {
    return Object.keys(variant).sort().map(g => `${g}=${variant[g]}`).join(",");
}

// Active (non-bypassed) flakes for a FlakeCombo node.
function activeComboFlakes(node) {
    return getComboFlakes(node).filter(f => !f.bypassed);
}

// Active (non-bypassed) presets for a FlakeModelCombo node. Bypass is keyed by
// preset name in node.properties._combo_bypassed.
function activeComboPresets(node) {
    const bypassed = new Set(node.properties?._combo_bypassed || []);
    return getComboPresets(node).filter(p => !bypassed.has(p));
}

// Number of prompts a generation will queue: the product of each
// FlakeCombo's active-flake count and each FlakeModelCombo's preset count.
// Returns 1 when there are combo nodes but a single combination, or 1 when
// there are no combo nodes at all (a normal single queue).
export function computeJobCount(graph) {
    const g = graph || app.graph;
    if (!g) return 1;
    let count = 1;
    for (const node of g.nodes) {
        if (FLAKE_COMBO_TYPES.includes(node.type)) {
            // Each active flake is one axis item; an "All Variants" flake expands
            // into one item per variant combination (#343), a "Multiple" one
            // into one per checked choice set (#406), each times its input images (#441).
            let axis = 0;
            for (const f of activeComboFlakes(node)) {
                const variants = f._all_variants ? Math.max(1, variantComboCount(f.name)) : expandEntryVariants(f).length;
                axis += variants * itemImageJobs(f, "").length;
            }
            if (axis > 0) count *= axis;
        } else if (MODEL_COMBO_TYPES.includes(node.type)) {
            const n = activeComboPresets(node).length;
            if (n > 0) count *= n;
        } else if (STACK_TYPES.includes(node.type)) {
            // Stack flakes flagged "All Variants" are extra axes (#343), and so
            // are a Flake Animate Stack's input images (#441).
            for (const { flake } of stackAllVariantsFlakes(node)) {
                count *= Math.max(1, variantComboCount(flake.name));
            }
            count *= Math.max(1, stackImageJobs(node).length);
        }
    }
    return count;
}

export function cartesianProduct(arrays) {
    if (arrays.length === 0) return [[]];
    const result = [];
    const head = arrays[0];
    const tail = cartesianProduct(arrays.slice(1));
    for (const h of head) {
        for (const t of tail) {
            result.push([h, ...t]);
        }
    }
    return result;
}

const _originalQueuePrompt = app.queuePrompt;

let _batchCleanup = null;

function getPromptIdFromEvent(e) {
    const detail = e.detail;
    return typeof detail === "string" ? detail : detail?.prompt_id;
}

// Live batch state per workflow (#373): running combo item per combo node id,
// results per FlakeGenerate node id, progress counters. Switching workflow tabs
// rebuilds every node and the batch's workflow may not be the active one when
// its events arrive, so nodes read this when they render instead of owning it.
const _batchStates = new WeakMap(); // workflow -> { generating, images, completed, total }
const _batchListeners = new Set();

// Nodes that sample a chain: they show previews and keep per-combination results.
const GENERATE_TYPES = ["FlakeGenerate", "FlakeGenerateAnimation"];

function activeWorkflow() {
    // Frontends without a workflow store have a single implicit workflow.
    return app.extensionManager?.workflow?.activeWorkflow ?? app;
}

function activeState() {
    return _batchStates.get(activeWorkflow());
}

// Combo node -> the running uid its last render highlighted.
const _shownUids = new WeakMap();

// Uid of the running item of a FlakeCombo / FlakeModelCombo node, if any.
// Called by the node's render.
export function getGeneratingUid(node) {
    const uid = activeState()?.generating[node.id];
    _shownUids.set(node, uid);
    return uid;
}

// { combo: {comboKey: img}, ad: {comboKey: img}, ab, last, lastKey } for a FlakeGenerate
// node, or null when no batch of this workflow produced results for it.
export function getBatchImages(node) {
    return activeState()?.images[node.id] || null;
}

export function getBatchProgress() {
    return activeState() || null;
}

// fn() runs after any batch state change; returns an unsubscribe function.
export function subscribeBatch(fn) {
    _batchListeners.add(fn);
    return () => _batchListeners.delete(fn);
}

// Re-renders only the combo nodes whose highlighted item changed.
function renderHighlights(state) {
    if (activeState() !== state) return;
    for (const n of app.graph.nodes) {
        if (_shownUids.has(n) && _shownUids.get(n) === state?.generating[n.id]) continue;
        if (FLAKE_COMBO_TYPES.includes(n.type)) n._combo_render?.();
        else if (MODEL_COMBO_TYPES.includes(n.type)) n._model_combo_render?.();
    }
}

// FlakeGenerate nodes read the active workflow's state themselves.
function renderResults() {
    for (const n of app.graph.nodes) {
        if (GENERATE_TYPES.includes(n.type)) n._flake_generate_render?.();
    }
    for (const fn of _batchListeners) fn();
}

// Latest per-step sampling preview (#374). Only one prompt executes at a time,
// so a single blob URL is kept; the previous one is revoked when replaced.
let _preview = null; // { url, nodeId, promptId }

function setPreview(next) {
    if (_preview) URL.revokeObjectURL(_preview.url);
    _preview = next;
    renderResults();
}

// { url, comboKey } while `node` is sampling, else null. comboKey is the
// running batch combination ("" for a plain run). Previews of a batch queued
// from another workflow are not shown.
export function getLivePreview(node) {
    if (!_preview || _preview.nodeId !== String(node.id)) return null;
    const idx = _activeBatch.promptIds.indexOf(_preview.promptId);
    if (idx >= 0 && _activeBatch.state !== activeState()) return null;
    return { url: _preview.url, comboKey: idx >= 0 ? _activeBatch.comboKeys[idx] : "" };
}

// FlakeGenerate shows its sampling previews itself; keep core from also drawing
// them on the node (#380). Registered at module load, i.e. before core's own
// listener (extensions load before the app adds its api handlers), so stopping
// propagation here hides the event from it.
_resolveComfyApi()?.addEventListener("b_preview_with_metadata", (e) => {
    const { blob, displayNodeId, jobId } = e.detail;
    if (!GENERATE_TYPES.includes(app.graph?.getNodeById(displayNodeId)?.type)) return;
    e.stopImmediatePropagation();
    setPreview({ url: URL.createObjectURL(blob), nodeId: String(displayNodeId), promptId: jobId });
});

// ComfyUI sends no sampling previews when the Live preview setting is
// "default" and the server runs without --preview-method (#380).
let _serverPreviewDefault = null;
function serverPreviewDefault() {
    _serverPreviewDefault ??= fetch("/flakes/preview_default").then(r => r.json()).then(d => d.method).catch(() => null);
    return _serverPreviewDefault;
}

// Set by setupBatchTracking while a batch runs; called with each prompt the
// batch loop queues.
let _onBatchPromptQueued = null;

// api.queuePrompt stays wrapped for the session: prompts with a FlakeGenerate
// get previews unless the user or the server chose a method, and batch prompts
// report their prompt_id (app.queuePrompt in current frontends returns void;
// api.queuePrompt returns the {prompt_id, number, node_errors} REST response).
function installQueuePromptWrapper(comfyApi) {
    const origQueuePrompt = comfyApi.queuePrompt.bind(comfyApi);
    comfyApi.queuePrompt = async function(number, data, opts) {
        // Snapshot the combo key/index now — they may already be cleared by
        // the time the awaits below resolve.
        const comboKey = _currentBatchComboKey;
        const comboIndex = _currentComboIndex;
        if ((!opts?.previewMethod || opts.previewMethod === "default")
            && Object.values(data?.output || {}).some(n => GENERATE_TYPES.includes(n?.class_type))
            && await serverPreviewDefault() === "none") {
            opts = { ...opts, previewMethod: "auto" };
        }
        const result = await origQueuePrompt(number, data, opts);
        if (result?.prompt_id && comboKey !== null) _onBatchPromptQueued?.(result.prompt_id, comboKey, comboIndex);
        return result;
    };
}

app.registerExtension({
    name: "comfyui-flakes.BatchState",
    setup() {
        const comfyApi = _resolveComfyApi();
        if (!comfyApi) return;
        installQueuePromptWrapper(comfyApi);
        // The node's result replaces its preview; a finished, interrupted or
        // superseded prompt drops it.
        comfyApi.addEventListener("executed", ({ detail }) => {
            if (_preview && String(detail.display_node || detail.node) === _preview.nodeId) setPreview(null);
        });
        const dropForPrompt = (e) => { if (_preview && getPromptIdFromEvent(e) === _preview.promptId) setPreview(null); };
        comfyApi.addEventListener("execution_success", dropForPrompt);
        comfyApi.addEventListener("execution_error", dropForPrompt);
        comfyApi.addEventListener("execution_interrupted", dropForPrompt);
        comfyApi.addEventListener("execution_start", (e) => { if (_preview && getPromptIdFromEvent(e) !== _preview.promptId) setPreview(null); });
        comfyApi.addEventListener("reconnected", () => { if (_preview) setPreview(null); });
    },
    // The workflow store switches the active workflow only after the new
    // graph's nodes are configured, so re-read the store once loading is complete.
    afterLoadGraph() {
        renderHighlights(activeState());
        renderResults();
    },
});

// ComfyUI exposes the websocket-event-emitting `api` singleton in slightly
// different places depending on frontend version. Try several spots so we
// don't silently lose batch tracking on a newer/older frontend.
function _resolveComfyApi() {
    return (
        window.comfyAPI?.api?.api
        || window.comfyAPI?.api
        || app.api
        || window.app?.api
        || null
    );
}

// Set by the outer batch loop before each combo's queue call; read by the
// patched api.queuePrompt so the captured prompt_id can be tagged with the
// right comboKey. Cleared between combos.
let _currentBatchComboKey = null;

// The combination index for the prompt currently being queued (#348). With a UI
// batch count > 1 each combination queues several prompts, so the prompt-id
// index no longer equals the combination index — this maps it back.
let _currentComboIndex = null;

// Active batch promptId → comboKey map, plus a live "running" pointer.
// Exported so the Generation Data overlay can show per-combo progress (#230).
const _activeBatch = { promptIds: [], comboKeys: [], runningPromptId: null, state: null };
export function getActiveBatch() {
    return _activeBatch;
}

function setupBatchTracking(state, combinations, comboKeys) {
    const comfyApi = _resolveComfyApi();
    if (!comfyApi || typeof comfyApi.addEventListener !== "function" || combinations.length === 0) {
        if (!comfyApi) {
            // eslint-disable-next-line no-console
            console.warn("[flakes] could not resolve ComfyUI api singleton; combo highlight/per-combo image storage disabled.");
        }
        return null;
    }
    const jobs = state.total;

    const promptIds = []; // queue order
    const comboIndexByPrompt = []; // parallel to promptIds: the combination index (#348)
    const done = new Set(); // finished prompt ids (events and /history backfill may both report one)
    let closed = false;
    _activeBatch.state = state;

    // Pending event buffer: websocket events for a prompt_id can arrive BEFORE
    // the REST response that gives us that prompt_id (different transport
    // channels). Buffer events we don't yet know about and replay them when
    // the patched api.queuePrompt finally registers the id. Without this, the
    // first generation in a batch consistently lost its execution_start /
    // executed events (#241).
    const pendingByPid = new Map(); // promptId -> [{type, e}]
    function applyEventForKnownPid(type, e, pid) {
        if (type === "execution_start") {
            setRunning(pid);
        } else if (type === "execution_success" || type === "execution_error" || type === "execution_interrupted") {
            markDone(pid);
        } else if (type === "executed") {
            handleOutput(pid, e.detail.display_node || e.detail.node, e.detail.output);
        }
    }
    function bufferOrApply(type, e) {
        const pid = type === "executed" ? (typeof e?.detail === "string" ? null : e?.detail?.prompt_id) : getPromptIdFromEvent(e);
        if (!pid) return;
        if (promptIds.includes(pid)) {
            applyEventForKnownPid(type, e, pid);
        } else {
            let arr = pendingByPid.get(pid);
            if (!arr) { arr = []; pendingByPid.set(pid, arr); }
            arr.push({ type, e });
        }
    }
    function flushPendingFor(pid) {
        const arr = pendingByPid.get(pid);
        if (!arr) return;
        pendingByPid.delete(pid);
        for (const { type, e } of arr) applyEventForKnownPid(type, e, pid);
    }

    // The api.queuePrompt wrapper reports each prompt_id the batch loop queues.
    _onBatchPromptQueued = (pid, comboKey, comboIndex) => {
        promptIds.push(pid);
        comboIndexByPrompt.push(comboIndex ?? 0);
        comboKeys.push(comboKey);
        _activeBatch.promptIds.push(pid);
        _activeBatch.comboKeys.push(comboKey);
        // Replay any websocket events that arrived before this push.
        flushPendingFor(pid);
    };

    function setRunning(pid) {
        const idx = promptIds.indexOf(pid);
        if (idx < 0) return;
        _activeBatch.runningPromptId = pid;
        state.generating = {};
        for (const item of combinations[comboIndexByPrompt[idx]]) {
            // stack_variant items have no node-level highlight (#343).
            if (item.type === "combo" || item.type === "model_combo") state.generating[item.node.id] = item.uid;
        }
        renderHighlights(state);
    }

    function markDone(pid) {
        if (!promptIds.includes(pid) || done.has(pid)) return;
        done.add(pid);
        state.completed = done.size;
        if (done.size >= jobs && _batchCleanup === cleanup) {
            _batchCleanup();
            _batchCleanup = null;
        }
    }

    // A FlakeGenerate output (from the `executed` event or /history): store
    // its images (or FlakeGenerateAnimation's videos) by the prompt's comboKey
    // under the producing node's id.
    function handleOutput(pid, nodeId, output) {
        const idx = promptIds.indexOf(pid);
        const images = output?.flake_images || output?.flake_videos;
        if (idx < 0 || !Array.isArray(images) || images.length === 0) return;
        const comboKey = comboKeys[idx] ?? "";
        const caption = output.flake_caption?.[0];
        const image = caption ? { ...images[0], caption } : images[0];
        const entry = state.images[nodeId] || (state.images[nodeId] = { combo: {}, ad: {}, ab: false, last: null, lastKey: null });
        entry.combo[comboKey] = image;
        entry.last = image;
        entry.lastKey = comboKey;
        const adImages = output.flake_images_ad;
        if (Array.isArray(output.adetailer_ab) && output.adetailer_ab.length > 0 && Array.isArray(adImages) && adImages.length > 0) {
            entry.ad[comboKey] = adImages[0];
            entry.ab = true;
        }
        renderResults();
    }

    // Websocket events are lost while the socket is down (background tabs,
    // reconnects): recover finished prompts from /history and the running one
    // from /queue. Prompts gone from both were deleted from the queue.
    async function backfill() {
        if (closed) return;
        const queue = await (await comfyApi.fetchApi("/queue")).json();
        if (closed) return;
        const queued = new Set([...queue.queue_running, ...queue.queue_pending].map(q => q[1]));
        for (const pid of promptIds.slice()) {
            if (done.has(pid) || queued.has(pid)) continue;
            const history = (await (await comfyApi.fetchApi(`/history/${pid}`)).json())[pid];
            if (closed) return;
            for (const [nodeId, output] of Object.entries(history?.outputs || {})) handleOutput(pid, nodeId, output);
            markDone(pid);
        }
        const running = queue.queue_running[0]?.[1];
        if (running && !closed) setRunning(running);
    }

    // All websocket lifecycle/output events route through bufferOrApply so
    // events arriving before their prompt_id is registered get queued.
    const onExecStart = (e) => bufferOrApply("execution_start", e);
    const onExecSuccess = (e) => bufferOrApply("execution_success", e);
    const onExecError = (e) => bufferOrApply("execution_error", e);
    const onExecInterrupted = (e) => bufferOrApply("execution_interrupted", e);
    const onExecuted = (e) => bufferOrApply("executed", e);
    const onVisibility = () => { if (document.visibilityState === "visible") backfill(); };

    comfyApi.addEventListener("execution_start", onExecStart);
    comfyApi.addEventListener("execution_success", onExecSuccess);
    comfyApi.addEventListener("execution_error", onExecError);
    comfyApi.addEventListener("execution_interrupted", onExecInterrupted);
    comfyApi.addEventListener("executed", onExecuted);
    comfyApi.addEventListener("reconnected", backfill);
    document.addEventListener("visibilitychange", onVisibility);

    function cleanup() {
        closed = true;
        comfyApi.removeEventListener("execution_start", onExecStart);
        comfyApi.removeEventListener("execution_success", onExecSuccess);
        comfyApi.removeEventListener("execution_error", onExecError);
        comfyApi.removeEventListener("execution_interrupted", onExecInterrupted);
        comfyApi.removeEventListener("executed", onExecuted);
        comfyApi.removeEventListener("reconnected", backfill);
        document.removeEventListener("visibilitychange", onVisibility);
        _onBatchPromptQueued = null;
        _currentBatchComboKey = null;
        _activeBatch.promptIds = [];
        _activeBatch.comboKeys = [];
        _activeBatch.runningPromptId = null;
        _activeBatch.state = null;
        pendingByPid.clear();
        state.generating = {};
        renderHighlights(state);
    }

    return { promptIds, cleanup };
}

// Order the combinatorial nodes (FlakeCombo + FlakeModelCombo) by their
// position in the pipeline connection chain (source -> sink), so the loop
// nesting follows the graph flow rather than node creation order. The
// FlakeModelCombo is the chain source (it loads the heavy checkpoint), so it
// lands first = outermost loop = reloaded least; each subsequent FlakeCombo
// toward FlakeGenerate becomes a deeper inner loop. Combo nodes not reachable
// from any FlakeGenerate (orphans) are appended in creation order so
// disconnected graphs keep working.
function orderedComboNodes() {
    const isCombo = (n) => FLAKE_COMBO_TYPES.includes(n.type) || MODEL_COMBO_TYPES.includes(n.type);
    const ordered = [];
    const seen = new Set();
    for (const gen of app.graph.nodes) {
        if (!GENERATE_TYPES.includes(gen.type)) continue;
        for (const { node } of generateChain(gen)) {
            if (isCombo(node) && !seen.has(node.id)) { seen.add(node.id); ordered.push(node); }
        }
    }
    for (const node of app.graph.nodes) {
        if (isCombo(node) && !seen.has(node.id)) { seen.add(node.id); ordered.push(node); }
    }
    // Exclude node-level bypassed (mode 4) / muted (mode 2) combo nodes entirely,
    // so they never contribute presets or raise "no active presets" errors (#304).
    return ordered.filter((n) => n.mode !== 4 && n.mode !== 2);
}

// A Generate node's upstream chain, source -> sink ([] when unconnected).
function generateChain(gen) {
    const fdInput = gen.inputs?.find(i => i.name === "flake_data" || i.name === "flake_anim_data");
    if (fdInput?.link == null) return [];
    const link = app.graph.links?.[fdInput.link];
    const startNode = link && app.graph.getNodeById(link.origin_id);
    if (!startNode) return [];
    try { return collectChain(startNode); } catch { return []; }
}

// Sets the combo / stack / preset widgets for one combination. Values are
// restored from the saved originals once queueing is done.
function applyCombination(combination) {
    for (const item of combination) {
        if (item.type === "combo") {
            const w = item.node.widgets?.find(w => w.name === "flakes_json");
            if (w) w.value = JSON.stringify([item.value]);
        } else if (item.type === "stack_variant") {
            // Set just this stack flake's variant for the combination;
            // other stack flakes (and other all-variants indices) are left
            // to their own items (#343).
            const w = item.node.widgets?.find(w => w.name === "flakes_json");
            if (w) {
                let arr; try { arr = JSON.parse(w.value || "[]"); } catch { arr = []; }
                if (arr[item.flakeIndex]) {
                    arr[item.flakeIndex] = { ...arr[item.flakeIndex], variant: item.value };
                    w.value = JSON.stringify(arr);
                }
            }
        } else if (item.type === "stack_image") {
            // This job's input image goes on its item only; the previous
            // job's is cleared (#441).
            const w = item.node.widgets?.find(w => w.name === "flakes_json");
            if (w) {
                let arr; try { arr = JSON.parse(w.value || "[]"); } catch { arr = []; }
                arr.forEach((e, i) => { e.image = i === item.flakeIndex ? item.value : undefined; });
                w.value = JSON.stringify(arr);
            }
        } else {
            const w = item.node.widgets?.find(w => w.name === "preset");
            if (w) w.value = item.value;
            // Per-instance preset overrides for this combination (#279).
            const ow = item.node.widgets?.find(w => w.name === "overrides_json");
            if (ow) {
                const ovr = (item.node.properties?._combo_overrides || [])[item.index] || {};
                ow.value = serializeModelOverrides(ovr);
            }
        }
    }
}

// FlakeAnimCaption job ({ preset, flakes, image }) for a Generate Animation
// node's chain as currently set, or null when no active flake directly in it
// has auto_caption (#484). Bypassed / muted chain nodes are skipped like the
// backend passes them through; the last start image in the chain wins.
async function captionJob(gen) {
    let preset = null, image = null, needed = false;
    const flakes = [];
    for (const { node } of generateChain(gen)) {
        if (node.mode === 2 || node.mode === 4) continue;
        if (!STACK_TYPES.includes(node.type) && !FLAKE_COMBO_TYPES.includes(node.type)) {
            preset = node.widgets?.find(w => w.name === "preset")?.value ?? null;
            continue;
        }
        flakes.push(node.widgets?.find(w => w.name === "flakes_json")?.value || "[]");
        for (const e of stackEntries(node)) {
            if (!e || e.bypassed) continue;
            if (e.image) image = e.image;
            if (e.inline || !e.name) continue;
            await fetchFlakeMeta(e.name);
            if (cachedAutoCaption(e.name)) needed = true;
        }
    }
    return needed ? { preset, flakes, image } : null;
}

// Queues one FlakeAnimCaption prompt and resolves with its captions (parallel
// to `jobs`), or null when it errors or is interrupted. Events that arrive
// before the REST response gives the prompt_id are replayed after it.
async function runCaptionPrompt(comfyApi, jobs) {
    const types = ["executed", "execution_error", "execution_interrupted"];
    let pid = null, finish;
    const early = [];
    const result = new Promise(resolve => { finish = resolve; });
    const onEvent = (e) => {
        if (pid === null) { early.push(e); return; }
        if (getPromptIdFromEvent(e) !== pid) return;
        if (e.type !== "executed") finish(null);
        else if (Array.isArray(e.detail.output?.flake_captions)) finish(e.detail.output.flake_captions);
    };
    for (const t of types) comfyApi.addEventListener(t, onEvent);
    try {
        const output = { flake_caption: { class_type: "FlakeAnimCaption", inputs: { jobs_json: JSON.stringify(jobs) } } };
        pid = (await comfyApi.queuePrompt(-1, { output })).prompt_id;
        early.splice(0).forEach(onEvent);
        return await result;
    } finally {
        for (const t of types) comfyApi.removeEventListener(t, onEvent);
    }
}

app.queuePrompt = async function(number, batchCount = 1) {
    const orderedNodes = orderedComboNodes();

    // Stack flakes flagged "All Variants" become extra combinatorial axes,
    // independent of any combo nodes (#343).
    const stackAxes = [];
    for (const node of app.graph.nodes) {
        if (!STACK_TYPES.includes(node.type) || node.mode === 2 || node.mode === 4) continue;
        for (const { flake, index } of stackAllVariantsFlakes(node)) {
            const meta = await fetchFlakeMeta(flake.name);
            const combos = enumerateVariantCombos(meta);
            if (combos.length <= 1) continue;
            stackAxes.push(combos.map((variant) => ({
                node, type: "stack_variant", flakeIndex: index, value: variant, vkey: `s${index}.${variantKey(variant)}`,
            })));
        }
        const imageJobs = stackImageJobs(node);
        if (imageJobs.length) {
            stackAxes.push(imageJobs.map(({ index, image, key }) => ({
                node, type: "stack_image", flakeIndex: index, value: image, vkey: key,
            })));
        }
    }

    if (orderedNodes.length === 0 && stackAxes.length === 0) {
        // A plain run supersedes this workflow's finished batch results.
        const state = activeState();
        if (state && state.completed >= state.total) _batchStates.delete(activeWorkflow());
        return _originalQueuePrompt.call(this, number, batchCount);
    }

    const optionsArrays = [];

    for (const node of orderedNodes) {
        if (FLAKE_COMBO_TYPES.includes(node.type)) {
            const flakes = getComboFlakes(node);
            if (flakes.length === 0) {
                window.alert(`${node.type} node has no flakes selected.`);
                return;
            }
            const activeFlakes = flakes.filter(flake => !flake.bypassed);
            if (activeFlakes.length === 0) {
                window.alert(`${node.type} node has no active (non-bypassed) flakes.`);
                return;
            }
            // Expand "All Variants" flakes into one axis item per variant combo,
            // "Multiple" choices into one per checked choice (#406); ordinary
            // flakes stay a single item. vkey keeps each combo's comboKey
            // unique (#343). Each splits again per input image (#441), the
            // job's entry carrying it as `image`.
            const comboItems = [];
            for (const flake of activeFlakes) {
                const uid = flake._uid;
                const variants = flake._all_variants ? enumerateVariantCombos(await fetchFlakeMeta(flake.name)) : expandEntryVariants(flake);
                const split = flake._all_variants || variants.length > 1;
                for (const variant of variants) {
                    const value = split ? { ...flake, variant } : flake;
                    for (const { image, key } of itemImageJobs(flake, split ? `${uid}.${variantKey(variant)}` : uid)) {
                        comboItems.push({ node, type: "combo", value: image ? { ...value, image } : value, uid, vkey: key });
                    }
                }
            }
            optionsArrays.push(comboItems);
        } else {
            const presets = getComboPresets(node);
            if (presets.length === 0) {
                window.alert(`${node.type} node has no presets selected.`);
                return;
            }
            const bypassed = new Set(node.properties?._combo_bypassed || []);
            const activePresets = presets
                .map((preset, i) => ({ preset, i }))
                .filter(({ preset }) => !bypassed.has(preset));
            if (activePresets.length === 0) {
                window.alert(`${node.type} node has no active (non-bypassed) presets.`);
                return;
            }
            // index picks the preset's overrides when its combination is queued.
            const uids = node.properties._combo_uids;
            optionsArrays.push(activePresets.map(({ preset, i }) => ({
                node,
                type: "model_combo",
                value: preset,
                index: i,
                uid: uids[i],
                vkey: uids[i],
            })));
        }
    }

    // Append the stack "All Variants" axes (#343).
    for (const axis of stackAxes) optionsArrays.push(axis);

    const combinations = cartesianProduct(optionsArrays);
    if (combinations.length === 0) {
        return _originalQueuePrompt.call(this, number, batchCount);
    }

    // Save original widget values (once per unique node)
    const nodeOriginals = new Map();
    for (const item of combinations[0]) {
        if (!nodeOriginals.has(item.node.id)) {
            if (item.type === "combo" || item.type === "stack_variant" || item.type === "stack_image") {
                // All mutate the node's flakes_json; save it once to restore later.
                const w = item.node.widgets?.find(w => w.name === "flakes_json");
                nodeOriginals.set(item.node.id, { node: item.node, widget: w, value: w?.value });
            } else {
                const w = item.node.widgets?.find(w => w.name === "preset");
                const ow = item.node.widgets?.find(w => w.name === "overrides_json");
                nodeOriginals.set(item.node.id, {
                    node: item.node, widget: w, value: w?.value,
                    overrideWidget: ow, overrideValue: ow?.value,
                });
            }
        }
    }

    const flakeGenerateNodes = app.graph.nodes.filter(n => GENERATE_TYPES.includes(n.type));
    const restoreWidgets = () => {
        for (const orig of nodeOriginals.values()) {
            if (orig.widget) orig.widget.value = orig.value;
            if (orig.overrideWidget) orig.overrideWidget.value = orig.overrideValue;
        }
    };

    // The UI batch count multiplies the combinations: each combination is run
    // `runsPerCombo` times in a row before moving to the next (#348). The seed
    // advances between runs via control_after_generate (randomize/increment).
    const runsPerCombo = Math.max(1, batchCount | 0);
    const totalJobs = combinations.length * runsPerCombo;

    // Auto-caption pre-pass (#484): per combination, the caption job of each
    // active Generate Animation node (null when its chain needs none). Every
    // queued video gets its own caption, so repeats are listed again.
    const animNodes = flakeGenerateNodes.filter(n => n.type === "FlakeGenerateAnimation" && n.mode !== 2 && n.mode !== 4);
    const captionJobs = [];
    try {
        for (const combination of combinations) {
            applyCombination(combination);
            captionJobs.push(await Promise.all(animNodes.map(captionJob)));
        }
    } finally {
        restoreWidgets();
    }
    const jobs = captionJobs.flatMap(js => Array(runsPerCombo).fill(js.filter(Boolean)).flat());
    for (const n of animNodes) {
        const w = n.widgets?.find(w => w.name === "caption");
        nodeOriginals.set(n.id, { node: n, widget: w, value: "" });
    }

    // Clean up any previous batch tracking
    if (_batchCleanup) {
        _batchCleanup();
        _batchCleanup = null;
    }

    // Fresh state for this workflow: the "[i/N] done" label restarts at "[0/N]"
    // (#227) and the previous batch's results are dropped.
    const workflow = activeWorkflow();
    const state = { generating: {}, images: {}, completed: 0, total: totalJobs, captioning: jobs.length > 0 };
    _batchStates.set(workflow, state);
    for (const n of flakeGenerateNodes) {
        n.properties._images_by_combo = {};
        n.properties._images_by_combo_ad = {};
        n.properties._adetailer_ab = false;
    }
    renderResults();

    let captions = null, captionIdx = 0;
    if (jobs.length) {
        try {
            captions = await runCaptionPrompt(_resolveComfyApi(), jobs);
        } catch (e) {
            console.error("[flakes] caption job failed", e);
        }
        state.captioning = false;
        // The videos are queued from the active graph, so a tab switch while
        // captioning would queue the wrong workflow.
        if (!captions || activeWorkflow() !== workflow) {
            if (_batchStates.get(workflow) === state) _batchStates.delete(workflow);
            renderResults();
            window.alert(captions ? "The workflow tab changed while captioning; no videos were queued."
                : "Auto-captioning failed or was cancelled; no videos were queued.");
            return;
        }
    }

    const comboKeys = []; // parallel to promptIds; comboKeys[idx] = key for promptIds[idx]
    const tracker = setupBatchTracking(state, combinations, comboKeys);
    if (tracker) {
        _batchCleanup = tracker.cleanup;
    }

    try {
        for (let comboIdx = 0; comboIdx < combinations.length; comboIdx++) {
            const combination = combinations[comboIdx];
            applyCombination(combination);
            const comboKey = combination
                .map(it => [it.node.id, it.vkey])
                .sort((a, b) => a[0] - b[0])
                .map(([id, vk]) => `${id}:${vk}`)
                .join("|");

            // Run this combination `runsPerCombo` times before the next (#348).
            // Tell our patched api.queuePrompt which comboKey/index to attach to
            // each captured prompt_id. app.queuePrompt in modern frontends returns
            // void; the prompt_id is captured inside the patch instead.
            for (let b = 0; b < runsPerCombo; b++) {
                if (captions) {
                    animNodes.forEach((n, i) => {
                        const w = n.widgets?.find(w => w.name === "caption");
                        const caption = captionJobs[comboIdx][i] ? captions[captionIdx++] : "";
                        if (w) w.value = caption;
                    });
                }
                if (tracker) { _currentBatchComboKey = comboKey; _currentComboIndex = comboIdx; }
                try {
                    await _originalQueuePrompt.call(this, number, 1);
                } finally {
                    if (tracker) { _currentBatchComboKey = null; _currentComboIndex = null; }
                }
            }
        }
    } finally {
        restoreWidgets();
    }
};