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

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

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

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

// 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 (node.type === "FlakeCombo") {
            const n = activeComboFlakes(node).length;
            if (n > 0) count *= n;
        } else if (node.type === "FlakeModelCombo") {
            const n = getComboPresets(node).length;
            if (n > 0) count *= n;
        }
    }
    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;
}

function clearAllHighlights() {
    for (const n of app.graph.nodes) {
        if (n.type === "FlakeCombo") {
            delete n._combo_generating_index;
            n._combo_render?.();
        }
        if (n.type === "FlakeModelCombo") {
            delete n._model_combo_generating_index;
            n._model_combo_render?.();
        }
    }
}

// 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;

// 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 };
export function getActiveBatch() {
    return _activeBatch;
}

function setupBatchTracking(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 promptIds = []; // insertion order = combo index
    let completedCount = 0;

    // 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") {
            const idx = promptIds.indexOf(pid);
            if (idx >= 0) { setComboHighlight(idx); _activeBatch.runningPromptId = pid; }
        } else if (type === "execution_success" || type === "execution_error" || type === "execution_interrupted") {
            const idx = promptIds.indexOf(pid);
            if (idx < 0) return;
            completedCount++;
            for (const n of app.graph.nodes) {
                if (n.type === "FlakeGenerate") {
                    n._batch_completed_count = completedCount;
                    n._batch_total_count = combinations.length;
                    n._batch_progress_render?.();
                }
            }
            maybeFinish();
        } else if (type === "executed") {
            handleExecuted(e);
        } else if (type === "progress") {
            // Progress events are handled by the overlay subscription; nothing
            // to replay batch-side. We only buffer the lifecycle events.
        }
    }
    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);
    }

    // Patch api.queuePrompt for the duration of this batch so we capture the
    // returned prompt_id. (app.queuePrompt in current frontend versions
    // returns void; the low-level api.queuePrompt still returns the
    // {prompt_id, number, node_errors} REST response.)
    const _origApiQueuePrompt = comfyApi.queuePrompt?.bind(comfyApi);
    if (_origApiQueuePrompt) {
        comfyApi.queuePrompt = async function(...args) {
            // Snapshot the combo key now — _currentBatchComboKey may already be
            // cleared by the time the await below resolves.
            const comboKeyForThisCall = _currentBatchComboKey;
            const result = await _origApiQueuePrompt(...args);
            if (result?.prompt_id && comboKeyForThisCall !== null) {
                promptIds.push(result.prompt_id);
                comboKeys.push(comboKeyForThisCall);
                _activeBatch.promptIds.push(result.prompt_id);
                _activeBatch.comboKeys.push(comboKeyForThisCall);
                // Replay any websocket events that arrived before this push.
                flushPendingFor(result.prompt_id);
            }
            return result;
        };
    }

    function setComboHighlight(comboIndex) {
        const combination = combinations[comboIndex];
        if (!combination) return;
        for (const item of combination) {
            if (item.type === "combo") {
                item.node._combo_generating_index = item.index;
                item.node._combo_render?.();
            } else {
                item.node._model_combo_generating_index = item.index;
                item.node._model_combo_render?.();
            }
        }
    }

    function maybeFinish() {
        if (completedCount >= combinations.length) {
            if (_batchCleanup) {
                _batchCleanup();
                _batchCleanup = null;
            }
        }
    }

    // The `executed` event carries the per-node output payload. Used to stamp
    // generated images onto FlakeGenerate nodes keyed by their comboKey.
    function handleExecuted(e) {
        const detail = e?.detail || {};
        const pid = typeof detail === "string" ? null : detail.prompt_id;
        const output = detail?.output;
        const nodeId = detail?.node;
        if (!pid || !output) return;
        const idx = promptIds.indexOf(pid);
        if (idx < 0) return;
        const comboKey = (comboKeys && comboKeys[idx]) ?? "";
        const images = output.flake_images;
        if (!Array.isArray(images) || images.length === 0) return;
        const candidates = nodeId
            ? [app.graph.getNodeById(nodeId)].filter(n => n && n.type === "FlakeGenerate")
            : app.graph.nodes.filter(n => n.type === "FlakeGenerate");
        for (const fg of candidates) {
            fg.properties = fg.properties || {};
            fg.properties._images_by_combo = fg.properties._images_by_combo || {};
            fg.properties._images_by_combo[comboKey] = images[0];
        }
    }

    // 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);

    comfyApi.addEventListener("execution_start", onExecStart);
    comfyApi.addEventListener("execution_success", onExecSuccess);
    comfyApi.addEventListener("execution_error", onExecError);
    comfyApi.addEventListener("execution_interrupted", onExecInterrupted);
    comfyApi.addEventListener("executed", onExecuted);

    const cleanup = () => {
        comfyApi.removeEventListener("execution_start", onExecStart);
        comfyApi.removeEventListener("execution_success", onExecSuccess);
        comfyApi.removeEventListener("execution_error", onExecError);
        comfyApi.removeEventListener("execution_interrupted", onExecInterrupted);
        comfyApi.removeEventListener("executed", onExecuted);
        if (_origApiQueuePrompt) {
            comfyApi.queuePrompt = _origApiQueuePrompt;
        }
        _currentBatchComboKey = null;
        _activeBatch.promptIds = [];
        _activeBatch.comboKeys = [];
        _activeBatch.runningPromptId = null;
        pendingByPid.clear();
        clearAllHighlights();
    };

    return { promptIds, cleanup };
}

app.queuePrompt = async function(number, batchCount = 1) {
    const comboNodes = app.graph.nodes.filter(n => n.type === "FlakeCombo");
    const modelComboNodes = app.graph.nodes.filter(n => n.type === "FlakeModelCombo");

    if (comboNodes.length === 0 && modelComboNodes.length === 0) {
        return _originalQueuePrompt.call(this, number, batchCount);
    }

    const optionsArrays = [];

    for (const node of comboNodes) {
        const flakes = getComboFlakes(node);
        if (flakes.length === 0) {
            window.alert("FlakeCombo node has no flakes selected.");
            return;
        }
        const activeFlakes = flakes
            .map((flake, i) => ({ flake, i }))
            .filter(({ flake }) => !flake.bypassed);
        if (activeFlakes.length === 0) {
            window.alert("FlakeCombo node has no active (non-bypassed) flakes.");
            return;
        }
        optionsArrays.push(activeFlakes.map(({ flake, i }) => ({
            node,
            type: "combo",
            value: flake,
            index: i,
        })));
    }

    for (const node of modelComboNodes) {
        const presets = getComboPresets(node);
        if (presets.length === 0) {
            window.alert("FlakeModelCombo node has no presets selected.");
            return;
        }
        optionsArrays.push(presets.map((preset, i) => ({
            node,
            type: "model_combo",
            value: preset,
            index: i,
        })));
    }

    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") {
                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");
                nodeOriginals.set(item.node.id, { node: item.node, widget: w, value: w?.value });
            }
        }
    }

    const flakeGenerateNodes = app.graph.nodes.filter(n => n.type === "FlakeGenerate");

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

    const comboKeys = []; // parallel to promptIds; comboKeys[idx] = key for promptIds[idx]
    const tracker = setupBatchTracking(combinations, comboKeys);
    if (tracker) {
        _batchCleanup = tracker.cleanup;
    }
    // Initialize progress counters on every FlakeGenerate node so the "[i/N] done"
    // label resets to "[0/N]" at batch start (#227).
    for (const n of flakeGenerateNodes) {
        n._batch_completed_count = 0;
        n._batch_total_count = combinations.length;
        n._batch_progress_render?.();
    }

    try {
        for (let comboIdx = 0; comboIdx < combinations.length; comboIdx++) {
            const combination = combinations[comboIdx];
            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 {
                    const w = item.node.widgets?.find(w => w.name === "preset");
                    if (w) w.value = item.value;
                }
            }
            const comboKey = combination
                .map(it => [it.node.id, it.index])
                .sort((a, b) => a[0] - b[0])
                .map(([id, idx]) => `${id}:${idx}`)
                .join("|");

            // Tell our patched api.queuePrompt which comboKey to attach to the
            // upcoming prompt_id. app.queuePrompt in modern frontends returns
            // void; the prompt_id is captured inside the patch instead.
            if (tracker) _currentBatchComboKey = comboKey;
            try {
                await _originalQueuePrompt.call(this, number, 1);
            } finally {
                if (tracker) _currentBatchComboKey = null;
            }
        }
    } finally {
        for (const orig of nodeOriginals.values()) {
            if (orig.widget) orig.widget.value = orig.value;
        }
        // If no tracker is available, clear highlights immediately as fallback
        if (!tracker) {
            clearAllHighlights();
        }
    }
};