import { LiteGraph, LGraph, LGraphCanvas, LGraphNode, type LGraphNodeConstructor, type LGraphNodeExecutable, type SerializedLGraph, type SerializedLGraphGroup, type SerializedLGraphNode, type SerializedLLink } from "@litegraph-ts/core";
import type { LConnectionKind, INodeSlot } from "@litegraph-ts/core";
import ComfyAPI from "$lib/api"
import { ComfyWidgets } from "$lib/widgets"
import defaultGraph from "$lib/defaultGraph"
import { getPngMetadata, importA1111 } from "$lib/pnginfo";
import EventEmitter from "events";
import type TypedEmitter from "typed-emitter";

// Import nodes
import * as basic from "@litegraph-ts/nodes-basic"
import * as nodes from "$lib/nodes/index"

LiteGraph.catch_exceptions = false;

if (typeof window !== "undefined") {
    // Load default visibility
    nodes.ComfyReroute.setDefaultTextVisibility(!!localStorage["Comfy.ComfyReroute.DefaultVisibility"]);
}

type QueueItem = { num: number, batchCount: number }

type ComfyAppEvents = {
    configured: (graph: LGraph) => void
    nodeAdded: (node: LGraphNode) => void
    nodeRemoved: (node: LGraphNode) => void
    nodeConnectionChanged: (kind: LConnectionKind, node: LGraphNode, slot: INodeSlot, targetNode: LGraphNode, targetSlot: INodeSlot) => void
    cleared: () => void
}

interface ComfyGraphNodeExecutable extends LGraphNodeExecutable {
    comfyClass: string
    isVirtualNode?: boolean;
    applyToGraph(workflow: SerializedLGraph<SerializedLGraphNode<LGraphNode>, SerializedLLink, SerializedLGraphGroup>): void;
}

export default class ComfyApp {
    api: ComfyAPI;
    canvasEl: HTMLCanvasElement | null = null;
    canvasCtx: CanvasRenderingContext2D | null = null;
    lGraph: LGraph | null = null;
    lCanvas: LGraphCanvas | null = null;
    dropZone: HTMLElement | null = null;
    nodeOutputs: Record<string, any> = {};
    eventBus: TypedEmitter<ComfyAppEvents> = new EventEmitter() as TypedEmitter<ComfyAppEvents>;

    private queueItems: QueueItem[] = [];
    private processingQueue: boolean = false;

    constructor() {
        this.api = new ComfyAPI();
    }

    async setup(): Promise<void> {
        this.addProcessMouseHandler();
        this.addProcessKeyHandler();

        this.canvasEl = document.getElementById("graph-canvas") as HTMLCanvasElement;
        this.lGraph = new LGraph();
        this.lCanvas = new LGraphCanvas(this.canvasEl, this.lGraph);
        this.canvasCtx = this.canvasEl.getContext("2d");

        this.addGraphLifecycleHooks();

        LiteGraph.release_link_on_empty_shows_menu = true;
        LiteGraph.alt_drag_do_clone_nodes = true;

        this.lGraph.start();

        // await this.#invokeExtensionsAsync("init");
        await this.registerNodes();

        // Load previous workflow
        let restored = false;
        try {
            const json = localStorage.getItem("workflow");
            if (json) {
                const workflow = JSON.parse(json);
                this.loadGraphData(workflow);
                restored = true;
            }
        } catch (err) {
            console.error("Error loading previous workflow", err);
        }

        // We failed to restore a workflow so load the default
        if (!restored) {
            this.loadGraphData();
        }

        // Save current workflow automatically
        setInterval(() => localStorage.setItem("workflow", JSON.stringify(this.lGraph.serialize())), 1000);

        // this.#addDrawNodeHandler();
        // this.#addDrawGroupsHandler();
        // this.#addApiUpdateHandlers();
        this.addDropHandler();
        // this.#addPasteHandler();
        // this.#addKeyboardHandler();

        // await this.#invokeExtensionsAsync("setup");

        // Ensure the canvas fills the window
        this.resizeCanvas();
        window.addEventListener("resize", this.resizeCanvas.bind(this));

        return Promise.resolve();
    }

    resizeCanvas() {
        // get current size of the canvas
        let rect = this.canvasEl.parentElement.getBoundingClientRect();

        // increase the actual size of our canvas
        this.canvasEl.width = rect.width * window.devicePixelRatio;
        this.canvasEl.height = rect.height * window.devicePixelRatio;

        // ensure all drawing operations are scaled
        this.canvasCtx.scale(window.devicePixelRatio, window.devicePixelRatio);

        // scale everything down using CSS
        this.canvasEl.style.width = rect.width + 'px';
        this.canvasEl.style.height = rect.height + 'px';
        this.lCanvas.draw(true, true);
    }

    private addProcessMouseHandler() {

    }

    private addProcessKeyHandler() {

    }

    private graphOnConfigure() {
        console.log("Configured");
        this.eventBus.emit("configured", this.lGraph);
    }

    private graphOnBeforeChange(graph: LGraph, info: any) {
        console.log("BeforeChange", info);
        this.eventBus.emit("beforeChange", graph, info);
    }

    private graphOnAfterChange(graph: LGraph, info: any) {
        console.log("AfterChange", info);
        this.eventBus.emit("afterChange", graph, info);
    }

    private graphOnNodeAdded(node: LGraphNode) {
        console.log("Added", node);
        this.eventBus.emit("nodeAdded", node);
    }

    private graphOnNodeRemoved(node: LGraphNode) {
        console.log("Removed", node);
        this.eventBus.emit("nodeRemoved", node);
    }

    private graphOnNodeConnectionChange(kind: LConnectionKind, node: LGraphNode, slot: INodeSlot, targetNode: LGraphNode, targetSlot: INodeSlot) {
        console.log("ConnectionChange", node);
        this.eventBus.emit("nodeConnectionChanged", kind, node, slot, targetNode, targetSlot);
    }

    private canvasOnClear() {
        console.log("CanvasClear");
        this.eventBus.emit("cleared");
    }

    private addGraphLifecycleHooks() {
        this.lGraph.onConfigure = this.graphOnConfigure.bind(this);
        this.lGraph.onBeforeChange = this.graphOnBeforeChange.bind(this);
        this.lGraph.onAfterChange = this.graphOnAfterChange.bind(this);
        this.lGraph.onNodeAdded = this.graphOnNodeAdded.bind(this);
        this.lGraph.onNodeRemoved = this.graphOnNodeRemoved.bind(this);
        this.lGraph.onNodeConnectionChange = this.graphOnNodeConnectionChange.bind(this);

        this.lCanvas.onClear = this.canvasOnClear.bind(this);
    }

    private async registerNodes() {
        const app = this;

        // Load node definitions from the backend
        const defs = await this.api.getNodeDefs();
        // await this.#invokeExtensionsAsync("addCustomNodeDefs", defs);

        // Generate list of known widgets
        const widgets = ComfyWidgets;
        // const widgets = Object.assign(
        // 	{},
        // 	ComfyWidgets,
        // 	...(await this.#invokeExtensionsAsync("getCustomWidgets")).filter(Boolean)
        // );

        // Register a node for each definition
        for (const nodeId in defs) {
            const nodeData = defs[nodeId];

            const ctor = class extends LGraphNode {
                constructor(title?: string) {
                    super(title);
                    this.type = nodeId; // XXX: workaround dependency in LGraphNode.addInput()
                    (this as any).comfyClass = nodeId;
                    var inputs = nodeData["input"]["required"];
                    if (nodeData["input"]["optional"] != undefined) {
                        inputs = Object.assign({}, nodeData["input"]["required"], nodeData["input"]["optional"])
                    }
                    const config = { minWidth: 1, minHeight: 1 };
                    for (const inputName in inputs) {
                        const inputData = inputs[inputName];
                        const type = inputData[0];

                        if (inputData[1]?.forceInput) {
                            this.addInput(inputName, type);
                        } else {
                            if (Array.isArray(type)) {
                                // Enums
                                Object.assign(config, widgets.COMBO(this, inputName, inputData, app) || {});
                            } else if (`${type}:${inputName}` in widgets) {
                                // Support custom widgets by Type:Name
                                Object.assign(config, widgets[`${type}:${inputName}`](this, inputName, inputData, app) || {});
                            } else if (type in widgets) {
                                // Standard type widgets
                                Object.assign(config, widgets[type](this, inputName, inputData, app) || {});
                            } else {
                                // Node connection inputs
                                this.addInput(inputName, type);
                            }
                        }
                    }

                    for (const o in nodeData["output"]) {
                        const output = nodeData["output"][o];
                        const outputName = nodeData["output_name"][o] || output;
                        this.addOutput(outputName, output);
                    }

                    const s = this.computeSize();
                    s[0] = Math.max(config.minWidth, s[0] * 1.5);
                    s[1] = Math.max(config.minHeight, s[1]);
                    this.size = s;
                    this.serialize_widgets = true;

                    // app.#invokeExtensionsAsync("nodeCreated", this);

                    return this;
                }
            }

            const node: LGraphNodeConstructor = {
                class: ctor,
                title: nodeData.name,
                type: nodeId,
                desc: `ComfyNode: ${nodeId}`
            }

            // this.#addNodeContextMenuHandler(node);
            // this.#addDrawBackgroundHandler(node, app);

            // await this.#invokeExtensionsAsync("beforeRegisterNodeDef", node, nodeData);
            LiteGraph.registerNodeType(node);
            node.category = nodeData.category;
        }

        // await this.#invokeExtensionsAsync("registerCustomNodes");
    }

    private showDropZone() {
        this.dropZone.style.display = "block";
    }

    private hideDropZone() {
        this.dropZone.style.display = "none";
    }

    private allowDrag(event: DragEvent) {
        if (event.dataTransfer.items?.length > 0) {
            event.dataTransfer.dropEffect = 'copy';
            this.showDropZone();
            event.preventDefault();
        }
    }

    private async handleDrop(event: DragEvent) {
        event.preventDefault();
        event.stopPropagation();
        this.hideDropZone();

        if (event.dataTransfer.files.length > 0) {
            await this.handleFile(event.dataTransfer.files[0]);
        }
    }

    private addDropHandler() {
        this.dropZone = document.getElementById("dropzone");

        window.addEventListener('dragenter', this.allowDrag.bind(this));
        this.dropZone.addEventListener('dragover', this.allowDrag.bind(this));
        this.dropZone.addEventListener('dragleave', this.hideDropZone.bind(this));
        this.dropZone.addEventListener('drop', this.handleDrop.bind(this));
    }

    /**
     * Populates the graph with the specified workflow data
     * @param {*} graphData A serialized graph object
     */
    loadGraphData(graphData: any = null) {
        this.clean();

        if (!graphData) {
            graphData = structuredClone(defaultGraph);
        }

        // Patch T2IAdapterLoader to ControlNetLoader since they are the same node now
        for (let n of graphData.nodes) {
            if (n.type == "T2IAdapterLoader") n.type = "ControlNetLoader";
        }

        this.lGraph.configure(graphData);

        for (const node of this.lGraph._nodes) {
            const size = node.computeSize();
            size[0] = Math.max(node.size[0], size[0]);
            size[1] = Math.max(node.size[1], size[1]);
            node.size = size;

            if (node.widgets) {
                // If you break something in the backend and want to patch workflows in the frontend
                // This is the place to do this
                for (let widget of node.widgets) {
                    if (node.type == "KSampler" || node.type == "KSamplerAdvanced") {
                        if (widget.name == "sampler_name") {
                            if (widget.value.constructor === String && widget.value.startsWith("sample_")) {
                                widget.value = widget.value.slice(7);
                            }
                        }
                    }
                }
            }

            // this.#invokeExtensions("loadedGraphNode", node);
        }
    }

    /**
     * Converts the current graph workflow for sending to the API
     * @returns The workflow and node links
     */
    async graphToPrompt(frontendState: Record<number, any[]> = {}) {
        const workflow = this.lGraph.serialize();

        const output = {};
        // Process nodes in order of execution
        for (const node of this.lGraph.computeExecutionOrder<ComfyGraphNodeExecutable>(false, null)) {
            const fromFrontend = frontendState[node.id];

            const n = workflow.nodes.find((n) => n.id === node.id);

            if (node.isVirtualNode || !node.comfyClass) {
                console.debug("Not serializing node: ", node.type)
                // Don't serialize frontend only nodes but let them make changes
                if (node.applyToGraph) {
                    node.applyToGraph(workflow);
                }
                continue;
            }

            if (node.mode === 2) {
                // Don't serialize muted nodes
                continue;
            }

            const inputs = {};
            const widgets = node.widgets;

            // Store all widget values
            if (widgets) {
                for (let i = 0; i < widgets.length; i++) {
                    const widget = widgets[i];
                    if (!widget.options || widget.options.serialize !== false) {
                        let value = widget.serializeValue ? await widget.serializeValue(n, i) : widget.value;
                        if (fromFrontend) {
                            value = fromFrontend[i].value;
                        }
                        inputs[widget.name] = value
                    }
                }
            }

            // Store all node links
            for (let i = 0; i < node.inputs.length; i++) {
                let parent: ComfyGraphNodeExecutable = node.getInputNode(i) as ComfyGraphNodeExecutable;
                if (parent) {
                    let link = node.getInputLink(i);
                    while (parent && parent.isVirtualNode) {
                        link = parent.getInputLink(link.origin_slot);
                        if (link) {
                            parent = parent.getInputNode(link.origin_slot) as ComfyGraphNodeExecutable;
                        } else {
                            parent = null;
                        }
                    }

                    if (link) {
                        inputs[node.inputs[i].name] = [String(link.origin_id), link.origin_slot];
                    }
                }
            }

            output[String(node.id)] = {
                inputs,
                class_type: node.comfyClass,
            };
        }

        // Remove inputs connected to removed nodes

        for (const o in output) {
            for (const i in output[o].inputs) {
                if (Array.isArray(output[o].inputs[i])
                    && output[o].inputs[i].length === 2
                    && !output[output[o].inputs[i][0]]) {
                    delete output[o].inputs[i];
                }
            }
        }

        return { workflow, output };
    }

    async queuePrompt(num: number, batchCount: number = 1, frontendState: Record<number, any[]> = {}) {
        this.queueItems.push({ num, batchCount });

        // Only have one action process the items so each one gets a unique seed correctly
        if (this.processingQueue) {
            return;
        }

        this.processingQueue = true;
        try {
            while (this.queueItems.length) {
                ({ num, batchCount } = this.queueItems.pop());
                console.log(`Queue get! ${num} ${batchCount}`);

                for (let i = 0; i < batchCount; i++) {
                    const p = await this.graphToPrompt(frontendState);

                    try {
                        await this.api.queuePrompt(num, p);
                    } catch (error) {
                        // this.ui.dialog.show(error.response || error.toString());
                        console.error(error.response || error.toString())
                        break;
                    }

                    for (const n of p.workflow.nodes) {
                        const node = this.lGraph.getNodeById(n.id);
                        if (node.widgets) {
                            for (const widget of node.widgets) {
                                // Allow widgets to run callbacks after a prompt has been queued
                                // e.g. random seed after every gen
                                // if (widget.afterQueued) {
                                // 	widget.afterQueued();
                                // }
                            }
                        }
                    }

                    this.lCanvas.draw(true, true);
                    // await this.ui.queue.update();
                }
            }
        } finally {
            console.log("Queue finished!");
            this.processingQueue = false;
        }
    }

    /**
     * Loads workflow data from the specified file
     */
    async handleFile(file: File) {
        if (file.type === "image/png") {
            const pngInfo = await getPngMetadata(file);
            if (pngInfo) {
                if (pngInfo.workflow) {
                    this.loadGraphData(JSON.parse(pngInfo.workflow));
                } else if (pngInfo.parameters) {
                    importA1111(this.lGraph, pngInfo.parameters, this.api);
                }
            }
        } else if (file.type === "application/json" || file.name.endsWith(".json")) {
            const reader = new FileReader();
            reader.onload = () => {
                this.loadGraphData(JSON.parse(reader.result as string));
            };
            reader.readAsText(file);
        }
    }

    // registerExtension(extension) {
    // 	if (!extension.name) {
    // 		throw new Error("Extensions must have a 'name' property.");
    // 	}
    // 	if (this.extensions.find((ext) => ext.name === extension.name)) {
    // 		throw new Error(`Extension named '${extension.name}' already registered.`);
    // 	}
    // 	this.extensions.push(extension);
    // }

    /**
     * Refresh combo list on whole nodes
     */
    async refreshComboInNodes() {
        const defs = await this.api.getNodeDefs();

        for (let nodeNum in this.lGraph._nodes) {
            const node = this.lGraph._nodes[nodeNum];

            const def = defs[node.type];

            for (const widgetNum in node.widgets) {
                const widget = node.widgets[widgetNum]

                if (widget.type == "combo" && def["input"]["required"][widget.name] !== undefined) {
                    widget.options.values = def["input"]["required"][widget.name][0];

                    if (!widget.options.values.includes(widget.value)) {
                        widget.value = widget.options.values[0];
                    }
                }
            }
        }
    }

    /**
     * Clean current state
     */
    clean() {
        this.nodeOutputs = {};
    }
}