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

const EXTENSION = "ethanfel.prompt_builder.accumulator_preview";
const NODE_NAME = "SxCPAccumulatorPreview";
const ANIM_PREVIEW_WIDGET = "$$comfy_animation_preview";
const ENTRY_ACTIONS = ["move up", "move down", "move top", "move bottom", "delete selected"];
const entryCache = new Map();

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

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

function resizeNode(node) {
    const size = node.computeSize?.();
    if (size) node.setSize?.(size);
    app.graph?.setDirtyCanvas(true, true);
}

function nodeKey(nodeOrId) {
    return String(typeof nodeOrId === "object" ? nodeOrId?.id : nodeOrId);
}

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 entryLabel(entry) {
    const index = entry?.index ?? "?";
    const id = entry?.id ? ` ${entry.id}` : "";
    const image = entry?.has_image ? " image" : " value";
    const shape = Array.isArray(entry?.shape) && entry.shape.length >= 2 ? ` ${entry.shape[1]}x${entry.shape[0]}` : "";
    const metadata = entry?.has_metadata ? " metadata" : "";
    return `#${index}${id}${image}${shape}${metadata}`.trim();
}

function setStatus(node, status) {
    if (!node._sxcpAccumulatorStatusWidget) return;
    node._sxcpAccumulatorStatusWidget.value = status || "no accumulator data";
    node.setDirtyCanvas?.(true, true);
}

function setEntries(node, entries, status = "") {
    entries = asArray(entries);
    entryCache.set(nodeKey(node), entries);
    node._sxcpAccumulatorEntries = entries;
    if (node._sxcpEntrySelectWidget) {
        const labels = entries.map(entryLabel);
        node._sxcpEntrySelectWidget.options.values = labels.length ? labels : ["no entries"];
        if (!labels.includes(node._sxcpEntrySelectWidget.value)) {
            node._sxcpEntrySelectWidget.value = labels[0] || "no entries";
        }
    }
    setStatus(node, status || `${entries.length} entries`);
    renderEntryStrip(node);
    resizeNode(node);
}

function selectedEntry(node) {
    const entries = entryCache.get(nodeKey(node)) || node._sxcpAccumulatorEntries || [];
    const selected = widget(node, "selected_entry")?.value || node._sxcpEntrySelectWidget?.value || "";
    const labels = entries.map(entryLabel);
    const index = labels.indexOf(selected);
    return index >= 0 ? entries[index] : entries[0];
}

function storeKey(node) {
    return String(widget(node, "store_key")?.value || node._sxcpResolvedStoreKey || "").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 removeAnimationPreviewWidget(node) {
    const widgetIndex = node.widgets?.findIndex((w) => w.name === ANIM_PREVIEW_WIDGET) ?? -1;
    if (widgetIndex < 0) return;
    node.widgets[widgetIndex].onRemove?.();
    node.widgets.splice(widgetIndex, 1);
}

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 loadImage(src) {
    return new Promise((resolve) => {
        const img = new Image();
        img.onload = () => resolve(img);
        img.onerror = () => resolve(null);
        img.src = src;
    });
}

function selectEntry(node, entry) {
    if (!node._sxcpEntrySelectWidget || !entry) return;
    const label = entryLabel(entry);
    node._sxcpEntrySelectWidget.value = label;
    node.setDirtyCanvas?.(true, true);
}

function entryStripRows(node) {
    const entries = entryCache.get(nodeKey(node)) || node._sxcpAccumulatorEntries || [];
    return entries.filter((entry) => entry?.has_image).slice(0, previewLimit(node));
}

function clearDragHighlights(root) {
    root?.querySelectorAll?.("[data-sxcp-drop]").forEach((row) => {
        row.style.borderColor = "rgba(255,255,255,0.18)";
        row.style.background = "rgba(255,255,255,0.04)";
        delete row.dataset.sxcpDrop;
    });
}

function dropTargetIndex(sourceEntry, targetEntry, event, targetElement) {
    const sourceIndex = Number(sourceEntry?.index || 0);
    const targetIndex = Number(targetEntry?.index || 0);
    if (!sourceIndex || !targetIndex) return 0;
    const rect = targetElement.getBoundingClientRect();
    const insertAfter = event.clientY > rect.top + rect.height / 2;
    let finalIndex = targetIndex + (insertAfter ? 1 : 0);
    if (sourceIndex < finalIndex) finalIndex -= 1;
    return Math.max(1, finalIndex);
}

function makeStripRow(node, entry, imageParams, rowNumber) {
    const row = document.createElement("div");
    row.draggable = true;
    row.dataset.index = String(entry.index || "");
    row.style.cssText = [
        "display:flex",
        "align-items:center",
        "gap:8px",
        "min-height:52px",
        "padding:4px",
        "border:1px solid rgba(255,255,255,0.18)",
        "border-radius:4px",
        "background:rgba(255,255,255,0.04)",
        "cursor:grab",
        "box-sizing:border-box",
        "user-select:none",
    ].join(";");
    row.title = "Drag to reorder accumulator images";

    const thumb = document.createElement("img");
    thumb.style.cssText = [
        "width:44px",
        "height:44px",
        "object-fit:cover",
        "border-radius:3px",
        "background:#111",
        "flex:0 0 auto",
    ].join(";");
    if (imageParams) thumb.src = imageUrl(imageParams);

    const label = document.createElement("div");
    label.textContent = `${rowNumber}. ${entryLabel(entry)}`;
    label.style.cssText = [
        "overflow:hidden",
        "text-overflow:ellipsis",
        "white-space:nowrap",
        "font:12px sans-serif",
        "color:#ddd",
        "flex:1 1 auto",
    ].join(";");

    row.append(thumb, label);
    row.addEventListener("click", () => selectEntry(node, entry));
    row.addEventListener("dragstart", (event) => {
        node._sxcpDraggedEntry = entry;
        row.style.opacity = "0.55";
        if (event.dataTransfer) {
            event.dataTransfer.effectAllowed = "move";
            event.dataTransfer.setData("text/plain", String(entry.id || entry.index || ""));
        }
    });
    row.addEventListener("dragend", () => {
        row.style.opacity = "1";
        clearDragHighlights(node._sxcpEntryStripRoot);
        node._sxcpDraggedEntry = null;
    });
    row.addEventListener("dragover", (event) => {
        event.preventDefault();
        if (!node._sxcpDraggedEntry || node._sxcpDraggedEntry === entry) return;
        clearDragHighlights(node._sxcpEntryStripRoot);
        row.dataset.sxcpDrop = "1";
        row.style.borderColor = "#7aa2ff";
        row.style.background = "rgba(122,162,255,0.18)";
    });
    row.addEventListener("drop", async (event) => {
        event.preventDefault();
        const sourceEntry = node._sxcpDraggedEntry;
        clearDragHighlights(node._sxcpEntryStripRoot);
        if (!sourceEntry || sourceEntry === entry) return;
        const targetIndex = dropTargetIndex(sourceEntry, entry, event, row);
        if (!targetIndex || targetIndex === Number(sourceEntry.index || 0)) return;
        await moveEntryToIndex(node, sourceEntry, targetIndex);
    });
    return row;
}

function renderEntryStrip(node) {
    const root = node._sxcpEntryStripRoot;
    if (!root) return;
    const entries = entryStripRows(node);
    const images = node._sxcpPreviewImageParams || [];
    root.replaceChildren();
    root.style.cssText = [
        "display:flex",
        "flex-direction:column",
        "gap:4px",
        "max-height:240px",
        "overflow:auto",
        "padding:4px",
        "border:1px solid rgba(255,255,255,0.12)",
        "border-radius:4px",
        "background:rgba(0,0,0,0.18)",
        "box-sizing:border-box",
    ].join(";");

    if (!entries.length) {
        const empty = document.createElement("div");
        empty.textContent = "no image entries";
        empty.style.cssText = "font:12px sans-serif;color:#aaa;padding:6px;";
        root.append(empty);
        node._sxcpEntryStripHeight = 36;
        return;
    }

    for (const [index, entry] of entries.entries()) {
        root.append(makeStripRow(node, entry, images[index], index + 1));
    }
    node._sxcpEntryStripHeight = Math.min(250, 12 + entries.length * 58);
}

function applyPreviewImages(node, images) {
    const normalized = asArray(images).filter(Boolean);
    node._sxcpPreviewImageParams = normalized;
    const key = nodeKey(node);
    const output = app.nodeOutputs?.[key] || {};
    if (app.nodeOutputs) {
        app.nodeOutputs[key] = {...output, images: normalized};
    }

    removeAnimationPreviewWidget(node);
    node.images = null;
    node.imgs = null;
    node.imageIndex = null;
    node.overIndex = null;
    node.pointerDown = null;

    if (!normalized.length) {
        renderEntryStrip(node);
        resizeNode(node);
        return;
    }

    Promise.all(normalized.map((params) => loadImage(imageUrl(params)))).then((loadedImages) => {
        if (app.nodeOutputs?.[key]?.images !== normalized) return;
        node.images = normalized;
        node.imgs = loadedImages.filter(Boolean);
        node.setSizeForImage?.();
        resizeNode(node);
    });
    renderEntryStrip(node);
    resizeNode(node);
}

function applyActionResult(node, data, status) {
    if (Object.prototype.hasOwnProperty.call(data || {}, "images")) {
        applyPreviewImages(node, data.images || []);
    }
    setEntries(node, data.entries || [], status || data.status || "");
}

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

async function refreshEntries(node) {
    const key = storeKey(node);
    if (!key) {
        alert("Set the same explicit store_key on the Accumulator and Accumulator Preview first.");
        return;
    }
    try {
        const data = await postJson("/sxcp/accumulator/list", actionPayload(node, {store_key: key}));
        applyActionResult(node, data, data.status || "");
    } catch (err) {
        console.error(`[${EXTENSION}] refresh failed`, err);
        alert(`Refresh failed: ${err}`);
    }
}

async function deleteSelected(node) {
    const key = storeKey(node);
    if (!key) {
        alert("Set the same explicit store_key on the Accumulator and Accumulator Preview first.");
        return;
    }
    const entry = selectedEntry(node);
    if (!entry) {
        alert("No accumulator entry selected.");
        return;
    }
    const label = entryLabel(entry);
    if (!confirm(`Delete accumulator entry ${label}?`)) 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,
        }));
        applyActionResult(node, data, `${data.status || ""}; deleted=${data.removed || 0}`);
    } catch (err) {
        console.error(`[${EXTENSION}] delete failed`, err);
        alert(`Delete failed: ${err}`);
    }
}

async function moveSelected(node, direction) {
    const key = storeKey(node);
    if (!key) {
        alert("Set the same explicit store_key on the Accumulator and Accumulator Preview first.");
        return;
    }
    const entry = selectedEntry(node);
    if (!entry) {
        alert("No accumulator entry selected.");
        return;
    }
    try {
        const data = await postJson("/sxcp/accumulator/move", actionPayload(node, {
            store_key: key,
            entry_id: entry.id || "",
            index: entry.id ? 0 : entry.index,
            direction,
        }));
        applyActionResult(node, data, `${data.status || ""}; moved=${data.moved ? "yes" : "no"}`);
    } catch (err) {
        console.error(`[${EXTENSION}] move failed`, err);
        alert(`Move 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,
        }));
        applyActionResult(node, data, `${data.status || ""}; moved=${data.moved ? "yes" : "no"}`);
    } catch (err) {
        console.error(`[${EXTENSION}] drag move failed`, err);
        alert(`Move failed: ${err}`);
    }
}

async function applyEntryAction(node) {
    const action = widget(node, "entry_action")?.value || node._sxcpEntryActionWidget?.value || "move up";
    if (action === "delete selected") {
        await deleteSelected(node);
        return;
    }
    const direction = action.replace(/^move\s+/, "");
    await moveSelected(node, direction);
}

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}));
        applyActionResult(node, data, `${data.status || ""}; cleared=${data.removed || 0}`);
    } catch (err) {
        console.error(`[${EXTENSION}] clear failed`, err);
        alert(`Clear failed: ${err}`);
    }
}

function setupNode(node) {
    hideWidget(widget(node, "delete_action"));
    hideWidget(widget(node, "delete_entry_id"));
    hideWidget(widget(node, "delete_index"));
    for (const legacyButton of [
        "Move Selected Top",
        "Move Selected Up",
        "Move Selected Down",
        "Move Selected Bottom",
        "Delete Selected Entry",
    ]) {
        hideWidget(widget(node, legacyButton));
    }

    if (!node._sxcpEntrySelectWidget) {
        node._sxcpEntrySelectWidget = node.addWidget("combo", "selected_entry", "no entries", () => {}, {values: ["no entries"]});
        node._sxcpEntrySelectWidget.serialize = false;
    }
    if (!node._sxcpEntryStripWidget && typeof node.addDOMWidget === "function") {
        node._sxcpEntryStripRoot = document.createElement("div");
        node._sxcpEntryStripWidget = node.addDOMWidget("entry_strip", "div", node._sxcpEntryStripRoot, {
            serialize: false,
            hideOnZoom: false,
            getMinHeight: () => node._sxcpEntryStripHeight || 36,
        });
        renderEntryStrip(node);
    }
    if (!node._sxcpEntryActionWidget) {
        node._sxcpEntryActionWidget = node.addWidget("combo", "entry_action", ENTRY_ACTIONS[0], () => {}, {values: ENTRY_ACTIONS});
        node._sxcpEntryActionWidget.serialize = false;
    }
    if (!node._sxcpAccumulatorStatusWidget) {
        node._sxcpAccumulatorStatusWidget = node.addWidget("text", "accumulator_status", "no accumulator data", () => {});
        node._sxcpAccumulatorStatusWidget.serialize = false;
    }
    if (!node._sxcpApplyEntryActionButton) {
        node._sxcpApplyEntryActionButton = node.addWidget("button", "Apply Entry Action", null, () => applyEntryAction(node));
    }
    if (!node._sxcpClearButton) {
        node._sxcpClearButton = node.addWidget("button", "Clear Accumulator", null, () => clearStore(node));
    }
    if (!node._sxcpRefreshButton) {
        node._sxcpRefreshButton = node.addWidget("button", "Refresh Entry List", null, () => refreshEntries(node));
    }
    resizeNode(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 || {};
            node._sxcpResolvedStoreKey = outputStoreKey(output);
            if (Object.prototype.hasOwnProperty.call(output, "images")) {
                applyPreviewImages(node, outputImages(output));
            }
            setEntries(node, outputEntries(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));
            return result;
        };
    },
});