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

const EXTENSION = "ethanfel.prompt_builder.accumulator_preview";
const NODE_NAME = "SxCPAccumulatorPreview";
const STYLE_ID = "sxcp-accumulator-preview-styles";

const MIN_CELL_W = 180;
const GAP = 6;
const PAD = 4;
const TOOLBAR_H = 28;
const CAPTION_H = 20;
const EMPTY_GRID_H = 84;
const MAX_GRID_H = 640;
const MIN_W = 560;
const MARGIN = 10;

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

function hideWidget(w) {
    if (!w) return;
    w.hidden = true;
    w.computeSize = () => [0, -4];
}

function isAccumulatorPreviewNode(node) {
    return node?.comfyClass === NODE_NAME || node?.type === NODE_NAME;
}

function getNodeById(id) {
    return app.graph?.getNodeById?.(Number(id)) || app.graph?._nodes_by_id?.[id] || app.graph?._nodes_by_id?.[Number(id)];
}

function asArray(value) {
    if (!value) return [];
    return Array.isArray(value) ? value : [value];
}

function outputStatus(output) {
    const status = output?.status;
    if (Array.isArray(status)) return status[0] || "";
    return status || "";
}

function outputStoreKey(output) {
    const key = output?.store_key;
    if (Array.isArray(key)) return key[0] || "";
    return key || "";
}

function outputEntries(output) {
    const entries = output?.entries;
    if (!entries) return [];
    if (Array.isArray(entries) && entries.length === 1 && Array.isArray(entries[0])) return entries[0];
    return asArray(entries);
}

function outputImages(output) {
    const images = output?.images;
    if (!images) return [];
    if (Array.isArray(images) && images.length === 1 && Array.isArray(images[0])) return images[0];
    return asArray(images).filter(Boolean);
}

function storeKey(node) {
    return String(widget(node, "store_key")?.value || node._sxapResolvedStoreKey || "").trim();
}

function previewLimit(node) {
    const value = Number(widget(node, "preview_limit")?.value ?? 64);
    return Number.isFinite(value) ? Math.max(1, Math.floor(value)) : 64;
}

function actionPayload(node, values = {}) {
    return {preview_limit: previewLimit(node), ...values};
}

function imageEntries(node) {
    const entries = node._sxapEntries || [];
    return entries.filter((entry) => entry?.has_image).slice(0, previewLimit(node));
}

function entryKey(entry) {
    const key = String(entry?.preview_key || "").trim();
    if (key) return `key:${key}`;
    const id = String(entry?.id || "").trim();
    if (id) return `id:${id}`;
    return `index:${entry?.index ?? ""}`;
}

function buildImageMap(entries, images, previous = new Map()) {
    const next = new Map(previous);
    const imageEntries = asArray(entries).filter((entry) => entry?.has_image);
    asArray(images).filter(Boolean).forEach((image, index) => {
        const directKey = String(image?.preview_key || "").trim();
        if (directKey) {
            next.set(`key:${directKey}`, image);
            return;
        }
        const entryId = String(image?.entry_id || "").trim();
        if (entryId) {
            next.set(`id:${entryId}`, image);
            return;
        }
        const entry = imageEntries[index];
        if (entry) next.set(entryKey(entry), image);
    });
    return next;
}

function imageParamsForEntry(node, entry, fallbackIndex) {
    const keyed = node._sxapImageByKey?.get(entryKey(entry));
    if (keyed) return keyed;
    return (node._sxapImages || [])[fallbackIndex];
}

function imageUrl(params) {
    const query = new URLSearchParams(params).toString();
    const preview = app.getPreviewFormatParam?.() || "";
    const rand = app.getRandParam?.() || `&rand=${Math.random()}`;
    return api.apiURL(`/view?${query}${preview}${rand}`);
}

function entryLabel(entry) {
    const parts = [`#${entry?.index ?? "?"}`];
    if (entry?.id) parts.push(entry.id);
    if (Array.isArray(entry?.shape) && entry.shape.length >= 2) parts.push(`${entry.shape[1]}x${entry.shape[0]}`);
    if (entry?.has_metadata) parts.push("metadata");
    return parts.join(" ");
}

function entryTitle(entry) {
    const value = entry?.value ? `\n${entry.value}` : "";
    return `${entryLabel(entry)}${value}`;
}

async function postJson(path, payload) {
    const response = await api.fetchApi(path, {
        method: "POST",
        headers: {"Content-Type": "application/json"},
        body: JSON.stringify(payload),
    });
    const data = await response.json();
    if (!response.ok) throw new Error(data?.error || response.statusText);
    return data;
}

function injectStyles() {
    if (document.getElementById(STYLE_ID)) return;
    const css = `
    .sxap-wrap { display:flex; flex-direction:column; gap:4px; box-sizing:border-box; height:100%; min-height:0; }
    .sxap-grid { display:flex; flex-wrap:wrap; gap:${GAP}px; align-content:flex-start; overflow-y:auto;
                 padding:${PAD}px; background:rgba(0,0,0,0.15); border-radius:4px;
                 flex:1 1 auto; min-height:0; box-sizing:border-box; }
    .sxap-grid.sxap-dragover { outline:2px dashed #6cf; outline-offset:-2px; }
    .sxap-empty { width:100%; padding:12px; text-align:center; font-size:12px; opacity:0.65; box-sizing:border-box; }
    .sxap-cell { position:relative; width:var(--sxap-cell-w, 240px); border:2px solid transparent;
                 border-radius:4px; overflow:hidden; background:#222; box-sizing:border-box;
                 transition:border-color .1s, background .1s; }
    .sxap-cell:hover { border-color:#555; }
    .sxap-cell.sxap-selected { border-color:#6cf; }
    .sxap-cell.sxap-drop { border-color:#fc6; border-style:dashed; }
    .sxap-thumb { width:100%; object-fit:contain; display:block; cursor:grab; background:#111; }
    .sxap-thumb:active { cursor:grabbing; }
    .sxap-badge { position:absolute; top:2px; left:2px; font-size:10px; background:rgba(0,0,0,0.65);
                  color:#fff; padding:0 4px; border-radius:3px; pointer-events:none; }
    .sxap-meta { position:absolute; top:3px; left:50%; transform:translateX(-50%); font-size:9px;
                 background:rgba(40,160,100,0.9); color:#fff; padding:0 3px; border-radius:3px;
                 pointer-events:none; }
    .sxap-del { position:absolute; top:2px; right:2px; width:18px; height:18px; border:none; border-radius:3px;
                background:rgba(180,30,30,0.88); color:#fff; font-size:11px; line-height:1;
                cursor:pointer; padding:0; }
    .sxap-caption { width:100%; height:${CAPTION_H}px; box-sizing:border-box;
                    overflow:hidden; text-overflow:ellipsis; white-space:nowrap; padding:2px 4px;
                    background:rgba(0,0,0,0.62); color:#fff; font-size:10px; line-height:16px; }
    .sxap-toolbar { display:flex; align-items:center; gap:6px; flex:0 0 auto; min-height:${TOOLBAR_H - 2}px; }
    .sxap-toolbar button { font-size:11px; padding:2px 8px; cursor:pointer; }
    .sxap-status { min-width:0; margin-left:auto; overflow:hidden; text-overflow:ellipsis; white-space:nowrap;
                   font-size:11px; opacity:0.75; }
    `;
    const style = document.createElement("style");
    style.id = STYLE_ID;
    style.textContent = css;
    document.head.appendChild(style);
}

function entryAspect(entry) {
    const shape = Array.isArray(entry?.shape) ? entry.shape : [];
    const height = Number(shape[0] || 0);
    const width = Number(shape[1] || 0);
    if (height > 0 && width > 0) return height / width;
    return 1;
}

function layoutMetrics(node) {
    const width = Math.max(node.size?.[0] || MIN_W, MIN_W);
    const inner = Math.max(MIN_CELL_W, width - 2 * MARGIN - 2 * PAD);
    const entries = imageEntries(node);
    const count = entries.length;
    const perRow = count <= 1 ? 1 : Math.max(1, Math.floor((inner + GAP) / (MIN_CELL_W + GAP)));
    const cellW = Math.max(MIN_CELL_W, Math.floor((inner - GAP * (perRow - 1)) / perRow));
    let gridH = EMPTY_GRID_H;
    if (count > 0) {
        gridH = 2 * PAD;
        for (let start = 0; start < count; start += perRow) {
            const row = entries.slice(start, start + perRow);
            const rowH = Math.max(...row.map((entry) => Math.round(cellW * entryAspect(entry)) + CAPTION_H));
            gridH += rowH;
            if (start + perRow < count) gridH += GAP;
        }
    }
    return {cellW, gridH};
}

function recomputeSize(node) {
    const {cellW, gridH} = layoutMetrics(node);
    node._sxapCellW = cellW;
    node._sxapWidgetH = 2 * MARGIN + TOOLBAR_H + 6 + Math.min(gridH, MAX_GRID_H);
}

function syncWidgetWidth(node) {
    if (node._sxapWidget) node._sxapWidget.width = node.size?.[0] || MIN_W;
}

function resizeToContent(node) {
    recomputeSize(node);
    const targetH = node.computeSize?.()[1] || node._sxapWidgetH || 140;
    const width = node.size?.[0] || MIN_W;
    const currentH = node.size?.[1] || 0;
    const previousAutoH = node._sxapAutoHeight || 0;
    const userSizedTaller = previousAutoH > 0 && currentH > previousAutoH + 6 && currentH > targetH;
    const nextH = userSizedTaller ? currentH : targetH;
    node._sxapAutoHeight = nextH;
    node.setSize?.([width, nextH]);
    syncWidgetWidth(node);
    node.setDirtyCanvas?.(true, true);
}

function suppressBuiltinPreview(node) {
    try {
        Object.defineProperty(node, "imgs", {
            configurable: true,
            get() { return undefined; },
            set() {},
        });
    } catch (_err) {
        // Best effort: if another extension made it non-configurable, the grid still works.
    }
}

function setWidgetValue(node, name, value) {
    const w = widget(node, name);
    if (!w) return;
    w.value = value;
    w.callback?.(value, app.canvas, node);
}

function setStatus(node, status) {
    node._sxapStatus = status || "";
    if (node._sxapStatusEl) node._sxapStatusEl.textContent = status || "";
    node.setDirtyCanvas?.(true, true);
}

function clearDropState(node) {
    node._sxapGridEl?.querySelectorAll?.(".sxap-drop").forEach((cell) => cell.classList.remove("sxap-drop"));
    node._sxapGridEl?.classList.remove("sxap-dragover");
}

function markSelected(node, selectedIndex) {
    node._sxapSelectedIndex = selectedIndex;
    node._sxapGridEl?.querySelectorAll?.(".sxap-cell").forEach((cell) => {
        cell.classList.toggle("sxap-selected", Number(cell.dataset.index || 0) === selectedIndex);
    });
}

function renderCell(node, entry, imageParams, displayIndex) {
    const cell = document.createElement("div");
    cell.className = "sxap-cell" + (entry.index === node._sxapSelectedIndex ? " sxap-selected" : "");
    cell.dataset.index = String(entry.index || "");
    cell.title = entryTitle(entry);
    const cellW = node._sxapCellW || MIN_CELL_W;
    cell.style.setProperty("--sxap-cell-w", `${cellW}px`);

    cell.ondragover = (event) => {
        if (!node._sxapDragEntry) return;
        event.preventDefault();
        event.stopPropagation();
        if (event.dataTransfer) event.dataTransfer.dropEffect = "move";
        cell.classList.add("sxap-drop");
    };
    cell.ondragleave = () => cell.classList.remove("sxap-drop");
    cell.ondrop = async (event) => {
        if (!node._sxapDragEntry) return;
        event.preventDefault();
        event.stopPropagation();
        cell.classList.remove("sxap-drop");
        const source = node._sxapDragEntry;
        node._sxapDragEntry = null;
        if (source.index === entry.index) return;
        await moveEntryToIndex(node, source, entry.index);
    };

    const thumb = document.createElement("img");
    thumb.className = "sxap-thumb";
    thumb.style.height = `${Math.max(48, Math.round(cellW * entryAspect(entry)))}px`;
    if (imageParams) thumb.src = imageUrl(imageParams);
    thumb.draggable = true;
    thumb.onclick = () => markSelected(node, entry.index);
    thumb.ondragstart = (event) => {
        node._sxapDragEntry = entry;
        if (event.dataTransfer) {
            event.dataTransfer.effectAllowed = "move";
            event.dataTransfer.setData("text/plain", String(entry.id || entry.index || ""));
        }
    };
    thumb.ondragend = () => {
        node._sxapDragEntry = null;
        clearDropState(node);
    };
    cell.appendChild(thumb);

    const badge = document.createElement("div");
    badge.className = "sxap-badge";
    badge.textContent = String(entry.index ?? displayIndex + 1);
    cell.appendChild(badge);

    if (entry.has_metadata) {
        const meta = document.createElement("div");
        meta.className = "sxap-meta";
        meta.textContent = "M";
        meta.title = "has metadata";
        cell.appendChild(meta);
    }

    const del = document.createElement("button");
    del.className = "sxap-del";
    del.textContent = "x";
    del.title = "Delete";
    del.onclick = async (event) => {
        event.stopPropagation();
        await deleteEntry(node, entry);
    };
    cell.appendChild(del);

    const caption = document.createElement("div");
    caption.className = "sxap-caption";
    caption.textContent = entry.id || entry.value || entryLabel(entry);
    cell.appendChild(caption);

    return cell;
}

function renderGrid(node) {
    const grid = node._sxapGridEl;
    if (!grid) return;
    const entries = imageEntries(node);
    grid.replaceChildren();

    if (!entries.length) {
        const empty = document.createElement("div");
        empty.className = "sxap-empty";
        empty.textContent = storeKey(node) ? "No accumulator images." : "Run once or set an explicit store_key.";
        grid.appendChild(empty);
    } else {
        entries.forEach((entry, index) => {
            grid.appendChild(renderCell(node, entry, imageParamsForEntry(node, entry, index), index));
        });
    }

    const total = (node._sxapEntries || []).filter((entry) => entry?.has_image).length;
    const count = entries.length;
    if (node._sxapStatusEl) {
        const prefix = total > count ? `${count}/${total} shown` : `${total} image${total === 1 ? "" : "s"}`;
        node._sxapStatusEl.textContent = node._sxapStatus ? `${prefix}; ${node._sxapStatus}` : prefix;
        node._sxapStatusEl.title = node._sxapStatusEl.textContent;
    }

    if (count !== node._sxapLastCount) {
        node._sxapLastCount = count;
        requestAnimationFrame(() => resizeToContent(node));
    } else {
        recomputeSize(node);
        syncWidgetWidth(node);
        node.setDirtyCanvas?.(true, true);
    }
}

function applyData(node, data, status = "") {
    if (Object.prototype.hasOwnProperty.call(data || {}, "store_key")) {
        node._sxapResolvedStoreKey = data.store_key || "";
    }
    if (Object.prototype.hasOwnProperty.call(data || {}, "entries")) {
        node._sxapEntries = asArray(data.entries);
    }
    if (Object.prototype.hasOwnProperty.call(data || {}, "images")) {
        node._sxapImages = asArray(data.images).filter(Boolean);
    }
    node._sxapImageByKey = buildImageMap(node._sxapEntries, node._sxapImages, node._sxapImageByKey);
    setStatus(node, status || data?.status || "");
    renderGrid(node);
}

async function refreshEntries(node) {
    const key = storeKey(node);
    if (!key) {
        setStatus(node, "no store_key yet");
        renderGrid(node);
        return;
    }
    try {
        const data = await postJson("/sxcp/accumulator/list", actionPayload(node, {store_key: key}));
        applyData(node, data, data.status || "");
    } catch (err) {
        console.error(`[${EXTENSION}] refresh failed`, err);
        alert(`Refresh failed: ${err}`);
    }
}

async function deleteEntry(node, entry) {
    const key = storeKey(node);
    if (!key) {
        alert("Set the same explicit store_key on the Accumulator and Accumulator Preview first.");
        return;
    }
    if (!entry) return;
    try {
        const data = await postJson("/sxcp/accumulator/delete", actionPayload(node, {
            store_key: key,
            entry_id: entry.id || "",
            index: entry.id ? 0 : entry.index,
            clear: false,
        }));
        applyData(node, data, `${data.status || ""}; deleted=${data.removed || 0}`);
    } catch (err) {
        console.error(`[${EXTENSION}] delete failed`, err);
        alert(`Delete failed: ${err}`);
    }
}

async function moveEntryToIndex(node, entry, targetIndex) {
    const key = storeKey(node);
    if (!key) {
        alert("Set the same explicit store_key on the Accumulator and Accumulator Preview first.");
        return;
    }
    if (!entry) return;
    try {
        const data = await postJson("/sxcp/accumulator/move", actionPayload(node, {
            store_key: key,
            entry_id: entry.id || "",
            index: entry.id ? 0 : entry.index,
            target_index: targetIndex,
        }));
        applyData(node, data, `${data.status || ""}; moved=${data.moved ? "yes" : "no"}`);
    } catch (err) {
        console.error(`[${EXTENSION}] drag move failed`, err);
        alert(`Move failed: ${err}`);
    }
}

async function clearStore(node) {
    const key = storeKey(node);
    if (!key) {
        alert("Set the same explicit store_key on the Accumulator and Accumulator Preview first.");
        return;
    }
    if (!confirm(`Clear all entries from accumulator "${key}"?`)) return;
    try {
        const data = await postJson("/sxcp/accumulator/delete", actionPayload(node, {store_key: key, clear: true}));
        applyData(node, data, `${data.status || ""}; cleared=${data.removed || 0}`);
    } catch (err) {
        console.error(`[${EXTENSION}] clear failed`, err);
        alert(`Clear failed: ${err}`);
    }
}

async function saveBatch(node) {
    const saveWidget = widget(node, "save_batch");
    if (!saveWidget) {
        alert("Missing save_batch widget.");
        return;
    }
    setWidgetValue(node, "finished", true);
    setWidgetValue(node, "save_batch", true);
    node.setDirtyCanvas?.(true, true);
    try {
        try {
            await app.queuePrompt(0, 1);
        } catch (_err) {
            await app.queuePrompt(0);
        }
    } catch (err) {
        console.error(`[${EXTENSION}] save queue failed`, err);
        alert(`Save failed: ${err}`);
    } finally {
        setWidgetValue(node, "save_batch", false);
        node.setDirtyCanvas?.(true, true);
    }
}

function hideInternalWidgets(node) {
    for (const name of [
        "delete_action",
        "delete_entry_id",
        "delete_index",
        "save_batch",
        "finished",
        "selected_entry",
        "entry_action",
        "accumulator_status",
    ]) {
        hideWidget(widget(node, name));
    }
    for (const legacyButton of [
        "Move Selected Top",
        "Move Selected Up",
        "Move Selected Down",
        "Move Selected Bottom",
        "Delete Selected Entry",
        "Apply Entry Action",
        "Clear Accumulator",
        "Refresh Entry List",
    ]) {
        hideWidget(widget(node, legacyButton));
    }
}

function setupNode(node) {
    injectStyles();
    suppressBuiltinPreview(node);
    hideInternalWidgets(node);

    if (!node._sxapGridEl && typeof node.addDOMWidget === "function") {
        const wrap = document.createElement("div");
        wrap.className = "sxap-wrap";

        const grid = document.createElement("div");
        grid.className = "sxap-grid";
        wrap.appendChild(grid);

        const toolbar = document.createElement("div");
        toolbar.className = "sxap-toolbar";

        const save = document.createElement("button");
        save.textContent = "Save";
        save.title = "Save current accumulator batch once";
        save.onclick = () => saveBatch(node);

        const refresh = document.createElement("button");
        refresh.textContent = "Refresh";
        refresh.onclick = () => refreshEntries(node);

        const clear = document.createElement("button");
        clear.textContent = "Clear";
        clear.onclick = () => clearStore(node);

        const status = document.createElement("span");
        status.className = "sxap-status";

        toolbar.append(save, refresh, clear, status);
        wrap.appendChild(toolbar);

        node._sxapGridEl = grid;
        node._sxapStatusEl = status;
        node._sxapEntries = node._sxapEntries || [];
        node._sxapImages = node._sxapImages || [];
        node._sxapImageByKey = node._sxapImageByKey || new Map();
        node._sxapLastCount = -1;
        recomputeSize(node);
        node._sxapWidget = node.addDOMWidget("accumulator_grid", "div", wrap, {
            serialize: false,
            getMinHeight: () => node._sxapWidgetH || 140,
        });
        renderGrid(node);
    }

    const onResize = node.onResize;
    if (!node._sxapResizeWrapped) {
        node.onResize = function () {
            const result = onResize?.apply(this, arguments);
            recomputeSize(node);
            syncWidgetWidth(node);
            return result;
        };
        node._sxapResizeWrapped = true;
    }

    syncWidgetWidth(node);
    if (!node._sxapInitialSizeSet) {
        node._sxapInitialSizeSet = true;
        node.setSize?.([Math.max(node.size?.[0] || 0, MIN_W), node.computeSize?.()[1] || node.size?.[1] || 180]);
    }
    renderGrid(node);
}

app.registerExtension({
    name: EXTENSION,

    async setup() {
        api.addEventListener("executed", ({detail}) => {
            const node = getNodeById(detail?.display_node ?? detail?.node);
            if (!isAccumulatorPreviewNode(node)) return;
            const output = detail?.output || {};
            const key = outputStoreKey(output);
            if (key) node._sxapResolvedStoreKey = key;
            applyData(node, {
                store_key: key,
                entries: outputEntries(output),
                images: outputImages(output),
                status: outputStatus(output),
            }, outputStatus(output));
        });
    },

    async beforeRegisterNodeDef(nodeType, nodeData) {
        if (nodeData.name !== NODE_NAME) return;

        const onNodeCreated = nodeType.prototype.onNodeCreated;
        nodeType.prototype.onNodeCreated = function () {
            const result = onNodeCreated?.apply(this, arguments);
            setupNode(this);
            return result;
        };

        const onConfigure = nodeType.prototype.onConfigure;
        nodeType.prototype.onConfigure = function () {
            const result = onConfigure?.apply(this, arguments);
            queueMicrotask(() => {
                setupNode(this);
                refreshEntries(this);
            });
            return result;
        };
    },
});