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

/* ------------------------------------------------------------------
   Extension: Comfy.GenerateTextNode
   Adds a read‑only multiline STRING widget to the "Generate Text (LLMToolkit)"
   node that displays the LLM response directly on the node.  The node still
   outputs the normal "context" so it can be connected to an external
   Display Text node as before.
-------------------------------------------------------------------*/

app.registerExtension({
    name: "Comfy.GenerateTextNode",
    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        // Only act on our LLMToolkitTextGenerator node
        if (nodeData.name !== "LLMToolkitTextGenerator") return;

        /* -----------------------------
           1.  Inject widget on creation
        ----------------------------- */
        const origOnNodeCreated = nodeType.prototype.onNodeCreated;
        nodeType.prototype.onNodeCreated = function () {
            // Invoke original handler first
            const ret = origOnNodeCreated ? origOnNodeCreated.apply(this, arguments) : undefined;

            // Build a unique widget name so each node instance has its own key
            const existingNodes = app.graph._nodes.filter(n => n.type === nodeData.name);
            const widgetName = `${nodeData.name}_response_${existingNodes.length}`;

            // Create read‑only STRING widget for displaying response
            const widgetInfo = ComfyWidgets.STRING(
                this,
                widgetName,
                ["STRING", {
                    default: "",
                    placeholder: "LLM response will appear here …",
                    multiline: true,
                }],
                app
            );

            // Make read‑only
            widgetInfo.widget.inputEl.readOnly = true;

            // Store reference for quick updates
            this.__llmResponseWidget = widgetInfo.widget || widgetInfo;

            // Ensure node size grows to fit another widget
            this.setSize(this.computeSize());

            return ret;
        };

        /* -----------------------------
           2.  Helper to update widget
        ----------------------------- */
        function updateResponseWidget(textContainer) {
            if (!this.__llmResponseWidget) return;

            // Extract final text string (textContainer may be array/string)
            let valueToSet = "";
            if (Array.isArray(textContainer)) {
                valueToSet = textContainer.join("\n");
            } else if (typeof textContainer === "string") {
                valueToSet = textContainer;
            }

            // Trim and write
            valueToSet = valueToSet.trim();
            this.__llmResponseWidget.value = valueToSet;
            if (this.__llmResponseWidget.inputEl) this.__llmResponseWidget.inputEl.value = valueToSet;
            app.graph.setDirtyCanvas(true);
        }

        /* -----------------------------
           3.  Hook into onExecuted to receive message
        ----------------------------- */
        const origOnExecuted = nodeType.prototype.onExecuted;
        nodeType.prototype.onExecuted = function (message) {
            if (origOnExecuted) origOnExecuted.apply(this, arguments);

            // message is typically { ui: { string: [...] }, result: (…) }
            let textData = undefined;
            if (message?.ui?.string) {
                textData = message.ui.string;
            } else if (message?.string) {
                textData = message.string;
            } else if (Array.isArray(message)) {
                textData = message;
            }
            updateResponseWidget.call(this, textData);
        };

        /* -----------------------------
           4. Restore saved value on load
        ----------------------------- */
        const origOnConfigure = nodeType.prototype.onConfigure;
        nodeType.prototype.onConfigure = function (config) {
            if (origOnConfigure) origOnConfigure.apply(this, arguments);

            if (!this.__llmResponseWidget || !config?.widgets_values?.length) return;

            const idx = this.widgets.findIndex(w => w === this.__llmResponseWidget);
            if (idx !== -1 && config.widgets_values.length > idx) {
                const savedVal = config.widgets_values[idx] || "";
                this.__llmResponseWidget.value = savedVal;
                if (this.__llmResponseWidget.inputEl) this.__llmResponseWidget.inputEl.value = savedVal;
            }
        };
    },
});