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

const EXTENSION = "ethanfel.prompt_builder.seed_control";
const NODE_NAME = "SxCPSeedControl";
const SEED_AXES = [
    "category",
    "subcategory",
    "content",
    "person",
    "scene",
    "pose",
    "role",
    "expression",
    "composition",
];

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

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

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

function randomSeed() {
    if (globalThis.crypto?.getRandomValues) {
        const values = new Uint32Array(1);
        globalThis.crypto.getRandomValues(values);
        return values[0];
    }
    return Math.floor(Math.random() * 0x100000000);
}

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

function materializeNode(node) {
    const changes = [];
    for (const axis of SEED_AXES) {
        const modeWidget = widget(node, `${axis}_seed_mode`);
        const seedWidget = widget(node, `${axis}_seed`);
        if (modeWidget?.value !== "random" || !seedWidget) continue;

        const seed = randomSeed();
        changes.push({
            node,
            modeWidget,
            previousMode: modeWidget.value,
            seedWidget,
            previousSeed: seedWidget.value,
            seed,
        });
        setWidgetValue(node, seedWidget, seed);
        setWidgetValue(node, modeWidget, "fixed");
    }
    if (changes.length) resizeNode(node);
    return changes;
}

function materializeAllForQueue() {
    const changes = [];
    for (const node of app.graph?._nodes || []) {
        if (isSeedControlNode(node)) changes.push(...materializeNode(node));
    }
    if (changes.length) app.graph?.setDirtyCanvas(true, true);
    return changes;
}

function restoreRandomModes(changes) {
    for (const change of changes) {
        setWidgetValue(change.node, change.modeWidget, change.previousMode);
    }
    if (changes.length) app.graph?.setDirtyCanvas(true, true);
}

function lockRandomSeeds(node) {
    const changes = materializeNode(node);
    if (!changes.length) {
        alert("No random seed modes to lock.");
        return;
    }
    resizeNode(node);
}

function setupNode(node) {
    node._sxcpSeedControlNode = true;
    if (!node._sxcpLockRandomSeedsButton) {
        node._sxcpLockRandomSeedsButton = node.addWidget("button", "Lock Random Seeds Now", null, () => lockRandomSeeds(node));
    }
    resizeNode(node);
}

app.registerExtension({
    name: EXTENSION,

    async setup() {
        if (app._sxcpSeedControlQueuePatched) return;
        const originalQueuePrompt = app.queuePrompt;
        if (!originalQueuePrompt) return;

        app._sxcpSeedControlQueuePatched = true;
        app.queuePrompt = async function () {
            const randomSeedChanges = materializeAllForQueue();
            try {
                return await originalQueuePrompt.apply(this, arguments);
            } finally {
                restoreRandomModes(randomSeedChanges);
            }
        };
    },

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