/**
 * ComfyUI basic data handling.
 * Copyright (C) 2025 StableLlama
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/**
 * Purpose: Manage dynamic input nodes in ComfyUI.
 * Inspired by cozy_ex_dynamic
 */

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

const TypeSlot = {
    Input: 1,
    Output: 2
};

const TypeSlotEvent = {
    Connect: true,
    Disconnect: false
};

app.registerExtension({
    name: 'Basic data handling: dynamic input',
    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        // Filter: Only process nodes with class names starting with "Basic data handling:"
        if (!nodeType.comfyClass.startsWith('Basic data handling:')) {
            return;
        }

        const combinedInputData = {
            ...nodeData?.input?.required ?? {},
            ...nodeData?.input?.optional ?? {}
        }
        const combinedInputDataOrder = [
            ...nodeData?.input_order?.required ?? [],
            ...nodeData?.input_order?.optional ?? []
        ];

        /** Array of (generic) dynamic inputs. */
        const dynamicInputs = [];
        /** Array of groups of (generic) dynamic inputs. */
        const dynamicInputGroups = [];
        for (const name of combinedInputDataOrder) {
            const dynamic = combinedInputData[name][1]?._dynamic;
            const dynamicGroup = combinedInputData[name][1]?._dynamicGroup ?? 0;

            if (dynamic) {
                let matcher;

                switch (dynamic) {
                    case 'number':
                        matcher = new RegExp(`^(${name.replace(/\d*$/, '')})(\\d+)$`);
                        break;

                    case 'letter':
                        matcher = new RegExp(`^()([a-zA-Z])$`);
                        break;

                    default:
                        continue;
                }
                const baseName = name.match(matcher)?.[1] ?? name;
                const dynamicType = combinedInputData[name][0];
                if (dynamicInputGroups[dynamicGroup] === undefined) {
                    dynamicInputGroups[dynamicGroup] = [];
                }
                dynamicInputGroups[dynamicGroup].push(
                    dynamicInputs.length
                )
                dynamicInputs.push({name, baseName, matcher, dynamic, dynamicType, dynamicGroup});
            }
        }
        if (dynamicInputs.length === 0) {
            return;
        }

        /**
         * Utility: Check if an input is dynamic.
         * @param {string} inputName - Name of the input to check.
         */
        const isDynamicInput = (inputName) =>
            dynamicInputs.some((di) => di.matcher.test(inputName));

        /**
         * Utility: Update inputs' slot indices after reordering.
         * @param {ComfyNode} node - The node to update.
         */
        const updateSlotIndices = (node) => {
            node.inputs.forEach((input, index) => {
                input.slot_index = index;
                if (input.isConnected) {
                    const link = node.graph._links.get(input.link);
                    if (link) {
                        link.target_slot = index;
                    } else {
                        console.error(`Input ${index} has an invalid link.`);
                    }
                }
            });
        };


        // Add helper method to insert input at a specific position
        nodeType.prototype.addInputAtPosition = function (name, type, position, isWidget, shape) {
            if (isWidget) {
                this.addWidget(type, name, '', ()=>{}, {});

                const GET_CONFIG = Symbol();
                const input = this.addInput(name, type, {
                    shape,
                    widget: {name, [GET_CONFIG]: () =>{}}
                  })
            } else {
                this.addInput(name, type, {shape}); // Add new input
            }
            const newInput = this.inputs.pop(); // Fetch the newly added input (last item)
            this.inputs.splice(position, 0, newInput); // Place it at the desired position
            updateSlotIndices(this); // Update indices
            return newInput;
        };

        // flag to prevent loops. It is ok to be "global" as the code is not
        // running in parallel.
        let isProcessingConnection = false;

        // Override onConnectionsChange: Handle connections for dynamic inputs
        const onConnectionsChange = nodeType.prototype.onConnectionsChange;
        nodeType.prototype.onConnectionsChange = function (type, slotIndex, isConnected, link, ioSlot) {
            const result = onConnectionsChange?.apply(this, arguments);

            if (type !== TypeSlot.Input || isProcessingConnection || !isDynamicInput(this.inputs[slotIndex].name)) {
                return result;
            }

            function getDynamicGroup(inputName) {
                // Find the dynamicGroup by matching the baseName with input name
                for (const di of dynamicInputs) {
                    if (inputName.startsWith(di.baseName)) {
                        return di.dynamicGroup;
                    }
                }
                return undefined;
            }

            isProcessingConnection = true;

            try {
                // Get dynamic input slots
                const dynamicSlots = [];
                const dynamicGroupCount = [];
                const dynamicGroupConnected = [];
                for (const [index, input] of this.inputs.entries()) {
                    const isDynamic = isDynamicInput(input.name);
                    if (isDynamic) {
                        const connected = input.isConnected;
                        const dynamicGroup = getDynamicGroup(input.name);
                        if (dynamicGroup in dynamicGroupCount) {
                            if (input.name.startsWith(dynamicInputs[dynamicInputGroups[dynamicGroup][0]].baseName)) {
                                dynamicGroupConnected[dynamicGroup][dynamicGroupCount[dynamicGroup]] ||= connected;
                                dynamicGroupCount[dynamicGroup]++;
                            }
                        } else {
                            dynamicGroupCount[dynamicGroup] = 1;
                            dynamicGroupConnected[dynamicGroup] = [connected];
                        }
                        dynamicSlots.push({
                            index,
                            name: input.name,
                            isWidget: input.widget !== undefined,
                            shape: input.shape,
                            connected,
                            isDynamic,
                            dynamicGroup,
                            dynamicGroupCount: dynamicGroupCount[dynamicGroup]
                        });

                        // sanity check to make sure every widget is in reality a widget. When loading a workflow this
                        // isn't the case so we must fix it ourselves.
                        if (this.widgets && !this.widgets.some((w) => w.name === input.name)) {
                            this.addWidget(input.type, input.name, '', ()=>{}, {});
                        }
                     }
                 }

                // Handle connection event
                if (isConnected === TypeSlotEvent.Connect) {
                    const hasEmptyDynamic = dynamicGroupConnected[0].some(dgc => !dgc);

                    if (!hasEmptyDynamic) {
                        // No empty slot - add a new one after the last dynamic input
                        const lastDynamicIdx = Math.max(...dynamicSlots.map((slot) => slot.index), -1);
                        let insertPosition = lastDynamicIdx + 1;
                        let inputInRange = true;

                        for (const groupMember of dynamicInputGroups[dynamicInputs[0].dynamicGroup]) {
                            const baseName = dynamicInputs[groupMember].baseName;
                            const dynamicType = dynamicInputs[groupMember].dynamicType;
                            let newName;
                            if (dynamicInputs[0].dynamic === 'letter') {
                                if (dynamicSlots.length >= 26) {
                                    inputInRange = false;
                                }
                                // For letter type, use the next letter in sequence
                                newName = String.fromCharCode(97 + dynamicSlots.length); // 97 is ASCII for 'a'
                            } else {
                                // For number type, use baseName + index as before
                                newName = `${baseName}${dynamicSlots.length}`;
                            }

                            if (inputInRange) {
                                // Insert the new empty input at the correct position
                                this.addInputAtPosition(newName, dynamicType, insertPosition++, dynamicSlots[groupMember].isWidget, dynamicSlots[groupMember].shape);
                                // Renumber inputs after addition
                                this.renumberDynamicInputs(baseName, dynamicInputs, dynamicInputs[0].dynamic);
                            }
                        }
                    }
                } else if (isConnected === TypeSlotEvent.Disconnect) {
                    let foundEmptyIndex = -1;

                    for (let idx = 0; idx < this.inputs.length; idx++) {
                        const input = this.inputs[idx];

                        if (!isDynamicInput(input.name)) {
                            continue;
                        }

                        if (!input.isConnected) { // Check if the input is empty
                            // remove empty input - but only when it's not the last one
                            const dynamicGroup = getDynamicGroup(input.name);
                            let isLast = true;
                            for (let i = idx + 1; i < this.inputs.length; i++) {
                                isLast &&= !this.inputs[i].name.startsWith(dynamicInputs[dynamicInputGroups[dynamicGroup][0]].baseName);
                            }
                            if (isLast) {
                                continue;
                            }

                            for (let i = idx + 1; i < this.inputs.length; i++) {
                                this.swapInputs(i - 1, i);
                            }
                            const lastIdx = this.inputs.length - 1;
                            if (this.inputs[lastIdx].widget !== undefined) {
                                const widgetIdx = this.widgets.findIndex((w) => w.name === this.inputs[lastIdx].widget.name)
                                this.widgets.splice(widgetIdx, 1);
                                this.widgets_values?.splice(widgetIdx, 1);
                            }
                            this.removeInput(lastIdx);
                        }
                    }

                    // Renumber dynamic inputs to ensure proper ordering
                    for (const groupMember of dynamicInputGroups[dynamicInputs[0].dynamicGroup]) {
                        const baseName = dynamicInputs[groupMember].baseName;
                        this.renumberDynamicInputs(baseName, dynamicInputs, dynamicInputs[0].dynamic);
                    }

                }

                this.setDirtyCanvas(true, true);
            } catch (e) {
                console.error(e);
                debugger;
                alert(e);
            } finally {
                isProcessingConnection = false;
            }

            return result;
        };

        const onConnectInput = nodeType.prototype.onConnectInput;
        nodeType.prototype.onConnectInput = function(inputIndex, outputType, outputSlot, outputNode, outputIndex) {
            const result = onRemoved?.apply(this, arguments) ?? true;

            if (this.inputs[inputIndex].isConnected) {
                const pre_isProcessingConnection = isProcessingConnection;
                isProcessingConnection = true;
                this.disconnectInput(inputIndex, true);
                isProcessingConnection = pre_isProcessingConnection;
            }
            return result;
        }

        const onRemoved = nodeType.prototype.onRemoved;
        nodeType.prototype.onRemoved = function () {
            const result = onRemoved?.apply(this, arguments);

            // When this is called, the input links are already removed - but
            // due to the implementation of the remove() method it might not
            // have worked with the dynamic inputs. So we need to fix it here.
            for (let i = this.inputs.length-1; i >= 0; i--) {
                if ( this.inputs[i].isConnected) {
                    this.disconnectInput(i, true);
                }
            }

            return result;
        }

        // Method to swap two inputs in the "this.inputs" array by their indices
        nodeType.prototype.swapInputs = function(indexA, indexB) {
            // Validate indices
            if (
                indexA < 0 || indexB < 0 ||
                indexA >= this.inputs.length ||
                indexB >= this.inputs.length ||
                indexA === indexB
            ) {
                console.error("Invalid input indices for swapping:", indexA, indexB);
                return;
            }

            // reflect the swap with the widgets
            if (this.inputs[indexA].widget !== undefined) {
                if (this.inputs[indexB].widget === undefined) {
                    console.error("Bad swap: input A is a widget but input B is not", indexA, indexB);
                }
                const widgetIdxA = this.widgets.findIndex((w) => w.name === this.inputs[indexA].widget.name);
                const widgetIdxB = this.widgets.findIndex((w) => w.name === this.inputs[indexB].widget.name);
                [this.widgets[widgetIdxA].y, this.widgets[widgetIdxB].y] = [this.widgets[widgetIdxB].y, this.widgets[widgetIdxA].y];
                [this.widgets[widgetIdxA].last_y, this.widgets[widgetIdxB].last_y] = [this.widgets[widgetIdxB].last_y, this.widgets[widgetIdxA].last_y];
                [this.widgets[widgetIdxA], this.widgets[widgetIdxB]] = [this.widgets[widgetIdxB], this.widgets[widgetIdxA]];
                if (this.widgets_values) {
                    [this.widgets_values[widgetIdxA], this.widgets_values[widgetIdxB]] = [this.widgets_values[widgetIdxB], this.widgets_values[widgetIdxA]];
                }
            }

            // Swap the inputs in the array
            [this.inputs[indexA].boundingRect, this.inputs[indexB].boundingRect] = [this.inputs[indexB].boundingRect, this.inputs[indexA].boundingRect];
            [this.inputs[indexA].pos, this.inputs[indexB].pos] = [this.inputs[indexB].pos, this.inputs[indexA].pos];
            [this.inputs[indexA], this.inputs[indexB]] = [this.inputs[indexB], this.inputs[indexA]];
            updateSlotIndices(this); // Refresh indices

            // Redraw the node to ensure the graph updates properly
            // -> not needed as the calling method must do it!
            this.setDirtyCanvas(true, true);
        };

        // Add method to safely renumber dynamic inputs without breaking connections
        nodeType.prototype.renumberDynamicInputs = function(baseName, dynamicInputs, dynamic) {
            // Get current dynamic inputs info
            const dynamicInputInfo = [];

            for (let i = 0; i < this.inputs.length; i++) {
                const input = this.inputs[i];
                const isDynamic = isDynamicInput(input.name);

                if (isDynamic && input.name.startsWith(baseName)) {
                    dynamicInputInfo.push({
                        index: i,
                        widgetIdx: input.widget !== undefined ? this.widgets.findIndex((w) => w.name === input.widget.name) : undefined,
                        name: input.name,
                        connected: input.isConnected
                    });
                }
            }

            // Just rename the inputs in place - don't remove/add to keep connections intact
            for (let i = 0; i < dynamicInputInfo.length; i++) {
                const info = dynamicInputInfo[i];
                const input = this.inputs[info.index];
                const newName = dynamic === "number" ? `${baseName}${i}` : String.fromCharCode(97 + i); // 97 is ASCII for 'a'

                if (input.widget !== undefined) {
                    const widgetIdx = info.widgetIdx;
                    input.widget.name = newName;
                    this.widgets[widgetIdx].name = newName;
                    this.widgets[widgetIdx].label = newName;
                }

                if (input.name !== newName) {
                    input.name = newName;
                    input.localized_name = newName;
                }
            }
        };
    }
});