import { app } from "../../../scripts/app.js";
 
function setupResolutionNode(node) {
    if (node._sata_res_initialized) return;
    if (node.comfyClass !== "Resolution_Machine" && node.comfyClass !== "Latent_Machine") return;

    // Find widgets
    const modelWidget = node.widgets?.find(w => w.name === "model");
    const dimensionWidget = node.widgets?.find(w => w.name === "dimension");
    const resolutionWidget = node.widgets?.find(w => w.name === "resolution");
    const widthWidget = node.widgets?.find(w => w.name === "width") || node.widgets?.find(w => w.name === "custom_width");
    const heightWidget = node.widgets?.find(w => w.name === "height") || node.widgets?.find(w => w.name === "custom_height");

    if (!modelWidget || !dimensionWidget || !resolutionWidget || !widthWidget || !heightWidget) {
        return;
    }

    node._sata_res_initialized = true;

    // Store resolution config
    let resolutionsConfig = {};

    // Fetch resolution config from backend
    async function loadConfig() {
        try {
            const resp = await fetch("/SATA_UtilityNode/resolutions_config");
            resolutionsConfig = await resp.json();
            refreshResolutions(false);
        } catch (err) {
            console.error("[Resolution_Machine] Failed to load config:", err);
        }
    }

    // --- Utility: refresh resolution dropdown when model/dimension changes ---
    function refreshResolutions(resetSelection = true) {
        const selectedModel = modelWidget.value;
        const selectedDimension = dimensionWidget.value;

        let available = [];

        if (resolutionsConfig.models && resolutionsConfig.models[selectedModel]) {
            const buckets = resolutionsConfig.models[selectedModel];
            buckets.forEach(bucket => {
                if (resolutionsConfig.resolutions && resolutionsConfig.resolutions[bucket] && resolutionsConfig.resolutions[bucket][selectedDimension]) {
                    available.push(...Object.keys(resolutionsConfig.resolutions[bucket][selectedDimension]));
                }
            });
        }

        // Always include "Custom" option
        if (!available.includes("Custom")) {
            available.push("Custom");
        }

        resolutionWidget.options.values = available;

        if (resetSelection && !available.includes(resolutionWidget.value)) {
            resolutionWidget.value = available[0] || "Custom";
        }

        checkCustomMode();
    }

    function checkCustomMode() {
        const isCustom = resolutionWidget.value === "Custom";
        if (isCustom) {
            widthWidget.type = "INT";
            heightWidget.type = "INT";
            delete widthWidget.computeSize;
            delete heightWidget.computeSize;
            widthWidget.disabled = false;
            heightWidget.disabled = false;
        } else {
            widthWidget.type = "hidden";
            heightWidget.type = "hidden";
            widthWidget.computeSize = () => [0, -4];
            heightWidget.computeSize = () => [0, -4];
            applyResolutionPreset();
        }

        if (node.computeSize && node.setSize) {
            node.setSize(node.computeSize());
        }
        if (app.graph) {
            if (app.graph.setDirtyCanvas) app.graph.setDirtyCanvas(true, true);
            if (app.graph.change) app.graph.change();
        }
    }

    // --- Utility: apply resolution preset to width/height ---
    function applyResolutionPreset() {
        const selectedModel = modelWidget.value;
        const selectedDimension = dimensionWidget.value;
        const selectedResolution = resolutionWidget.value;

        if (selectedResolution === "Custom") {
            return;
        }

        let resolutionData = null;

        if (resolutionsConfig.models && resolutionsConfig.models[selectedModel]) {
            const buckets = resolutionsConfig.models[selectedModel];
            for (const bucket of buckets) {
                if (resolutionsConfig.resolutions &&
                    resolutionsConfig.resolutions[bucket] &&
                    resolutionsConfig.resolutions[bucket][selectedDimension] &&
                    resolutionsConfig.resolutions[bucket][selectedDimension][selectedResolution]) {

                    resolutionData = resolutionsConfig.resolutions[bucket][selectedDimension][selectedResolution];
                    break;
                }
            }
        }

        if (resolutionData) {
            widthWidget.value = resolutionData.width;
            heightWidget.value = resolutionData.height;
        }
    }

    // --- Hook model change ---
    const oldModelCallback = modelWidget.callback;
    modelWidget.callback = function (value) {
        if (oldModelCallback) oldModelCallback.apply(this, arguments);
        refreshResolutions(true);
    };

    // --- Hook dimension change ---
    const oldDimensionCallback = dimensionWidget.callback;
    dimensionWidget.callback = function (value) {
        if (oldDimensionCallback) oldDimensionCallback.apply(this, arguments);
        refreshResolutions(true);
    };

    // --- Hook resolution change ---
    const oldResolutionCallback = resolutionWidget.callback;
    resolutionWidget.callback = function (value) {
        if (oldResolutionCallback) oldResolutionCallback.apply(this, arguments);
        checkCustomMode();
    };

    // --- Add Custom Button (once) ---
    if (!node.widgets.some(w => w.name === "Custom")) {
        node.addWidget("button", "Custom", null, () => {
            widthWidget.disabled = false;
            heightWidget.disabled = false;

            if (resolutionWidget.options && resolutionWidget.options.values) {
                if (!resolutionWidget.options.values.includes("Custom")) {
                    resolutionWidget.options.values.push("Custom");
                }
            }
            resolutionWidget.value = "Custom";

            checkCustomMode();
        });
    }

    loadConfig();
}

app.registerExtension({
    name: "SATA_UtilityNode.Resolution_Machine",

    async nodeCreated(node) {
        setupResolutionNode(node);
    },

    async loadedGraphNode(node) {
        setupResolutionNode(node);
    }
});