import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";
import { ImageCompareView } from "./compare.js?v=stability1";
import { readPreview } from "./preview_transport.js";
import { installComparisonRecovery } from "./result_recovery.js";
import {executionId,executionNode,graphNodeEntries} from "./workspace_flow.js";
import {upstreamNodes} from "./workbench_analysis.js";
import { appendComparison } from "./compare_history.js?v=0.5.8";
import { routeWheelOutsideImage } from "./wheel_routing.js";
import { T } from "./i18n.js";
import { bindWidgetToNodeWidth } from "./widget_layout.js?v=0.5.5";

function viewPath(image) {
    if (image?.type === "adachi_memory") return `/adachi-inpaint/compare/${encodeURIComponent(image.filename || "")}`;
    const type = ["input", "output", "temp"].includes(image?.type) ? image.type : "temp";
    return `/view?${new URLSearchParams({ filename: image?.filename || "", subfolder: image?.subfolder || "", type })}`;
}
const viewUrl = image => api.apiURL(viewPath(image));

function targets(node) {
    const graph=app.rootGraph||app.graph;
    const editors=graphNodeEntries(graph).filter(({node})=>node.type==='AdachiInlineMask' && node.adachiUseResult)
        .map(({node:editor,id})=>({id,label:`${editor.title||T('Inline Mask')} #${id}`}));
    const upstream=new Set([...upstreamNodes(node,'original')].filter(n=>n.type==='AdachiInlineMask').map(n=>executionId(graph,n)));
    return { editors, preferred: upstream.size === 1 ? [...upstream][0] : "" };
}

app.registerExtension({
    name: "Adachi.ImageCompare",
    setup() { installComparisonRecovery(api,()=>app.rootGraph||app.graph); },
    async beforeRegisterNodeDef(nodeType, nodeData) {
        if (nodeData.name !== "AdachiImageCompare") return;
        const created = nodeType.prototype.onNodeCreated;
        nodeType.prototype.onNodeCreated = function () {
            const value = created?.apply(this, arguments), node = this;
            const recovery=installComparisonRecovery(api,()=>app.rootGraph||app.graph);
            node.adachiTrackPrompt=(id,graph=app.rootGraph||app.graph)=>recovery.track(graph,id,executionId(graph,node));
            const compare = new ImageCompareView({
                resolveUrl: viewUrl,
                loadPreview: (image,options)=>readPreview(api.fetchApi.bind(api),viewPath(image),options),
                recover: ()=>recovery.recover(),
                waitingText: T("Waiting for the generated image…"),
                failedText: T("Generation failed; check the ComfyUI task history"),
                onGraphWheel: event => app.canvas?.processMouseWheel(event),
                getTargets: () => targets(node),
                onSelect: index => { node.properties.adachi_result_index = index; },
                async fetchResult(image, options) {
                    return readPreview(api.fetchApi.bind(api),viewPath(image),options);
                },
                async onAdopt(id, image, options={}) {
                    const graph=app.rootGraph||app.graph,target = executionNode(graph,id);
                    if (!target?.adachiUseResult) throw new Error(T("No inline mask node is available to receive the result"));
                    const blob = await readPreview(api.fetchApi.bind(api),viewPath(image),{...options,maxBytes:64*1024*1024});
                    options.signal?.throwIfAborted();
                    if(graph!==(app.rootGraph||app.graph)||target!==executionNode(graph,id))throw Error(T("The workflow changed; check connections again"));
                    if (blob.size > 64 * 1024 * 1024) throw new Error(T("Result image exceeds 64 MiB; downscale it first"));
                    await target.adachiUseResult(new File([blob], `adachi-result-${Date.now()}.png`, { type: "image/png" }), image.adachi_origin, options);
                    // Adoption starts a new round. Keep the accepted image visible
                    // as 1/1, then replace it with the next round's first outputs.
                    const data = { originals: [image], results: [image] };
                    node.properties.adachi_compare = data;
                    node.properties.adachi_result_index = 0;
                    node.properties.adachi_compare_new_round = true;
                    await compare.setData(data, 0);
                    app.graph?.change?.();
                },
            });
            node.adachiCompare = compare;
            const stopWheelRouting = routeWheelOutsideImage(compare.root, null, () => false);
            const widget = node.addDOMWidget("adachi_compare", "adachi_compare_view", compare.root, {
                serialize: false, hideOnZoom: false, getMinHeight: () => 420, getMaxHeight: () => 4096,
            });
            widget.serialize = false;
            bindWidgetToNodeWidth(node, widget, compare.root);
            const executed = node.onExecuted;
            node.onExecuted = function (output) {
                const result = executed?.apply(this, arguments);
                if (!output?.adachi_originals?.length || !output?.adachi_results?.length) return result;
                const previous = node.properties.adachi_compare_new_round
                    ? appendComparison() : appendComparison(node.properties.adachi_compare);
                const data = appendComparison(previous, {
                    originals: output.adachi_originals,
                    results: output.adachi_results.map(image => output.adachi_origin?.[0]
                        ? {...image, adachi_origin:output.adachi_origin[0]} : image),
                });
                node.properties.adachi_compare = data;
                node.properties.adachi_latest_result = data.results[data.results.length - 1];
                compare.setLatestResult(node.properties.adachi_latest_result);
                node.properties.adachi_compare_new_round = false;
                node.properties.adachi_result_index = data.results.length > previous.results.length
                    ? data.results.length - 1 : Math.max(0, Math.min(data.results.length - 1, node.properties.adachi_result_index || 0));
                compare.setData(data, node.properties.adachi_result_index);
                return result;
            };
            const configured = node.onConfigure;
            node.onConfigure = function () {
                const result = configured?.apply(this, arguments);
                if (node.properties.adachi_compare) {
                    node.properties.adachi_compare = appendComparison(node.properties.adachi_compare);
                    compare.setLatestResult(node.properties.adachi_latest_result || node.properties.adachi_compare.results.at(-1));
                    node.properties.adachi_result_index = Math.max(0, Math.min(node.properties.adachi_compare.results.length - 1, node.properties.adachi_result_index || 0));
                    compare.setData(node.properties.adachi_compare, node.properties.adachi_result_index);
                }
                return result;
            };
            const removed = node.onRemoved;
            node.onRemoved = function () { stopWheelRouting(); compare.destroy(); return removed?.apply(this, arguments); };
            node.setSize([Math.max(550, node.size[0]), Math.max(550, node.size[1])]);
            return value;
        };
    },
});