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

/* =============================================================================
   Prompt Database — ComfyUI Web Extension
   ========================================================================== */

console.log("[PromptDB] Extension loading...");

// ─── API helpers ────────────────────────────────────────────────────────────
const API = {
  async getCategories() {
    const r = await api.fetchApi("/prompt_db/categories");
    const data = await r.json();
    return data.categories || [];
  },
  async getPrompts(category) {
    const url = category
      ? `/prompt_db/prompts?category=${encodeURIComponent(category)}`
      : "/prompt_db/prompts";
    const r = await api.fetchApi(url);
    const data = await r.json();
    return data.prompts || [];
  },
};

function getNodeClass(node) {
  return node.comfyClass || node.type || node.constructor?.type || "";
}

// ─── Setup PromptDBManager ──────────────────────────────────────────────────
function setupManagerNode(node) {
  if (node._pdb_manager_setup) return;
  node._pdb_manager_setup = true;
  node.addWidget("button", "📝 Open Prompt Manager", null, () => {
    window.open("/prompt_db/manager", "_blank");
  });
  node.setSize?.(node.computeSize?.() || [260, 58]);
  app.graph?.setDirtyCanvas?.(true);
}

// ─── Setup PromptSelector ───────────────────────────────────────────────────
function setupSelectorNode(node) {
  if (node._pdb_selector_setup) return;
  node._pdb_selector_setup = true;

  const catWidget = node.widgets?.find((w) => w.name === "category");
  const promptWidget = node.widgets?.find((w) => w.name === "prompt_name");
  if (!catWidget || !promptWidget) return;

  // Strip 📖 marker before sending to Python backend
  promptWidget.serializeValue = () =>
    (promptWidget.value || "").replace(/\s*📖$/, "");

  // ── Refresh categories + prompts from DB ──
  async function refreshFromDB() {
    try {
      const categories = await API.getCategories();
      if (categories.length > 0) {
        catWidget.options.values = categories;
        if (!categories.includes(catWidget.value)) catWidget.value = categories[0];
      }
      const cat = catWidget.value;
      if (cat && cat !== "(no categories)") {
        const prompts = await API.getPrompts(cat);
        const names = prompts.map((p) => {
          const hasLoras = p.loras && p.loras.length > 0;
          return hasLoras ? `${p.name} 📖` : p.name;
        });
        if (names.length > 0) {
          promptWidget.options.values = names;
          if (!names.includes(promptWidget.value)) promptWidget.value = names[0];
        } else {
          promptWidget.options.values = ["(no prompts)"];
          promptWidget.value = "(no prompts)";
        }
      }
      await updateLoraDisplay();
      app.graph?.setDirtyCanvas?.(true);
    } catch (e) {
      console.error("[PromptDB] Refresh failed:", e);
    }
  }

  // ── LoRA display widgets (read-only info) ──
  async function updateLoraDisplay() {
    if (node.widgets) {
      node.widgets = node.widgets.filter((w) => !w._pdb_lora_display);
    }
    const rawName = promptWidget.value || "";
    const cleanName = rawName.replace(/\s*📖$/, "");
    const cat = catWidget.value;
    if (!cleanName || cleanName === "(no prompts)") return;

    try {
      const prompts = await API.getPrompts(cat);
      const prompt = prompts.find((p) => p.name === cleanName);
      if (!prompt || !prompt.loras || prompt.loras.length === 0) return;

      for (const lora of prompt.loras) {
        const displayName = (lora.name || "").replace(/\.safetensors$/, "");
        const strength = lora.strength != null ? lora.strength : 1.0;
        const w = node.addWidget(
          "text",
          `🔗 ${displayName}`,
          `${strength.toFixed(2)}`,
          () => {},
          { serialize: false }
        );
        w._pdb_lora_display = true;
        w.disabled = true;
      }
    } catch (e) {
      console.error("[PromptDB] Failed to load LoRA info:", e);
    }
    const computed = node.computeSize?.();
    if (computed) node.setSize?.(computed);
    app.graph?.setDirtyCanvas?.(true);
  }

  // ── Add Refresh button ──
  node.addWidget("button", "🔄 Refresh", null, () => refreshFromDB());

  // ── Category change callback ──
  const origCatCallback = catWidget.callback;
  catWidget.callback = async function (value) {
    if (origCatCallback) origCatCallback.call(this, value);
    try {
      const prompts = await API.getPrompts(value);
      const names = prompts.map((p) => {
        const hasLoras = p.loras && p.loras.length > 0;
        return hasLoras ? `${p.name} 📖` : p.name;
      });
      if (names.length === 0) {
        promptWidget.options.values = ["(no prompts)"];
        promptWidget.value = "(no prompts)";
      } else {
        promptWidget.options.values = names;
        if (!names.includes(promptWidget.value)) promptWidget.value = names[0];
      }
      await updateLoraDisplay();
      app.graph?.setDirtyCanvas?.(true);
    } catch (e) {
      console.error("[PromptDB] Failed to fetch prompts:", e);
    }
  };

  // ── Prompt change callback ──
  const origPromptCallback = promptWidget.callback;
  promptWidget.callback = async function (value) {
    if (origPromptCallback) origPromptCallback.call(this, value);
    await updateLoraDisplay();
  };

  refreshFromDB();
}

// ─── Setup PromptCombiner ───────────────────────────────────────────────────
const DEFAULT_WIDTH = 500;

function setupCombinerNode(node) {
  if (node._pdb_combiner_setup) return;
  node._pdb_combiner_setup = true;
  console.log("[PromptDB] Setting up PromptCombiner node", node.id);

  const selectionsWidget = node.widgets?.find(
    (w) => w.name === "prompt_selections"
  );
  if (selectionsWidget) {
    selectionsWidget.type = "converted-widget";
    selectionsWidget.computeSize = () => [0, -4];
    selectionsWidget.serializeValue = () => selectionsWidget.value;
  }

  let userWidth = Math.max(node.size?.[0] || DEFAULT_WIDTH, DEFAULT_WIDTH);

  const origOnResize = node.onResize;
  node.onResize = function (size) {
    if (origOnResize) origOnResize.apply(this, arguments);
    if (size && size[0]) userWidth = size[0];
  };

  let rebuildPending = false;
  let rebuildTimer = null;

  function getMaxSelectionIndex() {
    let max = 0;
    for (const inp of node.inputs || []) {
      const m = inp.name.match(/^selection_(\d+)$/);
      if (m) max = Math.max(max, parseInt(m[1]));
    }
    return max;
  }

  function syncSelectionsJSON() {
    if (!selectionsWidget) return;
    const obj = {};
    if (node.widgets) {
      for (const w of node.widgets) {
        if (w._pdb_dynamic && w._pdb_inputName) {
          // Skip muted entries
          if (w._pdb_muted) {
            obj[w._pdb_inputName] = "";
          } else {
            const val = (w.value || "").replace(/\s*📖$/, "");
            obj[w._pdb_inputName] = val;
          }
        }
      }
    }
    selectionsWidget.value = JSON.stringify(obj);
  }

  function getSourceCategory(slotIndex) {
    const input = node.inputs?.[slotIndex];
    if (!input || input.link == null) return null;
    const linkInfo = app.graph.links[input.link];
    if (!linkInfo) return null;
    const sourceNode = app.graph.getNodeById(linkInfo.origin_id);
    if (!sourceNode) return null;
    const catWidget = sourceNode.widgets?.find((w) => w.name === "category");
    return catWidget?.value || null;
  }

  async function fullRebuild() {
    // Remove dynamic widgets
    if (node.widgets) {
      node.widgets = node.widgets.filter(
        (w) => !w._pdb_dynamic && !w._pdb_lora_display
      );
    }

    // Cleanup disconnected slots
    if (node.inputs) {
      for (let i = node.inputs.length - 1; i >= 0; i--) {
        const inp = node.inputs[i];
        if (inp.name.startsWith("selection_") && inp.link == null) {
          node.removeInput(i);
        }
      }
    }

    const nextIdx = getMaxSelectionIndex() + 1;
    node.addInput(`selection_${nextIdx}`, "DB_SELECTION");

    let savedSelections = {};
    try {
      savedSelections = JSON.parse(selectionsWidget?.value || "{}");
    } catch (e) {}

    for (let i = 0; i < (node.inputs?.length || 0); i++) {
      const inp = node.inputs[i];
      if (!inp.name.startsWith("selection_")) continue;
      if (inp.link == null) continue;

      const category = getSourceCategory(i);
      if (!category) continue;

      let allPrompts = [];
      let promptNames = [];
      try {
        allPrompts = await API.getPrompts(category);
        promptNames = allPrompts.map((p) => {
          const hasLoras = p.loras && p.loras.length > 0;
          return hasLoras ? `${p.name} 📖` : p.name;
        });
      } catch (e) {
        console.error("[PromptDB] fetch failed for", category, e);
      }
      if (promptNames.length === 0) promptNames = ["(no prompts)"];

      const saved = savedSelections[inp.name];
      let defaultVal = promptNames[0];
      if (saved) {
        const match = promptNames.find(
          (n) =>
            n === saved ||
            n === `${saved} 📖` ||
            n.replace(/\s*📖$/, "") === saved
        );
        if (match) defaultVal = match;
      }

      // Check saved mute state
      const muteKey = `pdb_mute_${node.id}_${inp.name}`;
      const isMuted = localStorage.getItem(muteKey) === "1";

      // ── Custom single-line widget: ● 📂 Category   ◀ PromptName ▶ ──
      const closureCat = category;
      const closureNames = promptNames;
      const closureInpName = inp.name;
      const widget = {
        type: "pdb_combo",
        name: `📂 ${category}`,
        value: defaultVal,
        _pdb_dynamic: true,
        _pdb_inputName: closureInpName,
        _pdb_category: closureCat,
        _pdb_allPrompts: allPrompts,
        _pdb_muted: isMuted,
        _pdb_values: closureNames,
        options: { serialize: false },
        last_y: 0,
        computeSize: function () {
          return [node.size[0], 20];
        },
        // ── Custom draw ──
        draw: function (ctx, nodeRef, widgetWidth, posY, H) {
          this.last_y = posY;
          const margin = 15;
          const innerW = widgetWidth - margin * 2;

          // Background bar (standard LiteGraph color)
          const bgColor = LiteGraph.WIDGET_BGCOLOR || "#29293e";
          ctx.fillStyle = bgColor;
          ctx.beginPath();
          ctx.roundRect
            ? ctx.roundRect(margin, posY, innerW, H, H * 0.3)
            : ctx.rect(margin, posY, innerW, H);
          ctx.fill();

          // Toggle circle (left side)
          const toggleX = margin + 14;
          const toggleY = posY + H * 0.5;
          const r = 5;
          ctx.beginPath();
          ctx.arc(toggleX, toggleY, r, 0, Math.PI * 2);
          ctx.fillStyle = this._pdb_muted ? "#555" : "#7cacf8";
          ctx.fill();
          if (!this._pdb_muted) {
            ctx.beginPath();
            ctx.arc(toggleX - 1, toggleY - 1, r * 0.35, 0, Math.PI * 2);
            ctx.fillStyle = "rgba(255,255,255,0.25)";
            ctx.fill();
          }

          // Category name (after circle)
          ctx.fillStyle = this._pdb_muted ? "#777" : "#ddd";
          ctx.font = "12px sans-serif";
          ctx.textAlign = "left";
          ctx.textBaseline = "middle";
          const catText = `📂 ${closureCat}`;
          ctx.fillText(catText, margin + 28, posY + H * 0.5);
          const catTextW = ctx.measureText(catText).width;
          const catEnd = margin + 28 + catTextW + 8; // where category text ends

          // Right side layout: ◀ PromptName ▶ (right-aligned)
          const rightEdge = margin + innerW - 10;
          const displayVal = (this.value || "").replace(/\s*📖$/, "");

          // Measure prompt name and truncate if overlaps category
          ctx.font = "12px sans-serif";
          const maxTextW = rightEdge - catEnd - 36; // space for both arrows + gaps
          let valText = displayVal;
          if (ctx.measureText(valText).width > maxTextW) {
            while (ctx.measureText(valText + "…").width > maxTextW && valText.length > 1) {
              valText = valText.slice(0, -1);
            }
            valText = valText + "…";
          }

          const textW = ctx.measureText(valText).width;
          // Positions: ...◀ [text] ▶
          const arrowRight = rightEdge;
          const textRight = arrowRight - 14;           // text right edge (before ▶)
          const textLeft = textRight - textW;           // text left edge
          const arrowLeftX = textLeft - 14;             // ◀ position

          // Cache for mouse handler
          this._pdb_arrowLeft = arrowLeftX;
          this._pdb_catEnd = catEnd;

          // ◀ button
          ctx.fillStyle = "#aaa";
          ctx.textAlign = "center";
          ctx.font = "10px sans-serif";
          ctx.fillText("◀", arrowLeftX, posY + H * 0.5);

          // Prompt name (right-aligned)
          ctx.fillStyle = this._pdb_muted ? "#777" : "#ddd";
          ctx.font = "12px sans-serif";
          ctx.textAlign = "right";
          ctx.fillText(valText, textRight, posY + H * 0.5);

          // ▶ button
          ctx.fillStyle = "#aaa";
          ctx.textAlign = "center";
          ctx.font = "10px sans-serif";
          ctx.fillText("▶", arrowRight, posY + H * 0.5);

          ctx.textBaseline = "alphabetic";
        },
        // ── Custom mouse ──
        mouse: function (event, pos, nodeRef) {
          if (event.type !== "pointerdown" && event.type !== "mousedown")
            return false;

          const x = pos[0];
          const widgetWidth = nodeRef.size[0];
          const margin = 15;
          const innerW = widgetWidth - margin * 2;
          const rightEdge = margin + innerW - 10;
          const arrowLeftX = this._pdb_arrowLeft || (margin + 120);

          // Toggle area (left 45px — circle click)
          if (x < margin + 45) {
            this._pdb_muted = !this._pdb_muted;
            localStorage.setItem(muteKey, this._pdb_muted ? "1" : "0");
            syncSelectionsJSON();
            updateCombinerLoras();
            app.graph?.setDirtyCanvas?.(true);
            return true;
          }

          // ◀ button area
          if (x > arrowLeftX - 10 && x < arrowLeftX + 10) {
            const idx = this._pdb_values.indexOf(this.value);
            const newIdx = (idx - 1 + this._pdb_values.length) % this._pdb_values.length;
            this.value = this._pdb_values[newIdx];
            syncSelectionsJSON();
            updateCombinerLoras();
            app.graph?.setDirtyCanvas?.(true);
            return true;
          }

          // ▶ button area
          if (x > rightEdge - 10) {
            const idx = this._pdb_values.indexOf(this.value);
            const newIdx = (idx + 1) % this._pdb_values.length;
            this.value = this._pdb_values[newIdx];
            syncSelectionsJSON();
            updateCombinerLoras();
            app.graph?.setDirtyCanvas?.(true);
            return true;
          }

          // Text area click — open dropdown
          if (x > arrowLeftX + 10 && x < rightEdge - 10) {
            const that = this;
            const canvas = app.canvas;
            const ref_window = canvas.getCanvasWindow();
            // Clean display names for the menu
            const cleanValues = this._pdb_values.map(v => v.replace(/\s*📖$/, ""));
            new LiteGraph.ContextMenu(
              this._pdb_values,
              {
                scale: Math.max(1, canvas.ds?.scale || 1),
                event: event,
                className: "dark",
                callback: function (v) {
                  if (v) {
                    that.value = v;
                    syncSelectionsJSON();
                    updateCombinerLoras();
                    app.graph?.setDirtyCanvas?.(true);
                  }
                },
              },
              ref_window
            );
            return true;
          }

          return false;
        },
      };
      node.widgets.push(widget);

      // Add LoRA info strips for the currently selected prompt
      const cleanVal = defaultVal.replace(/\s*📖$/, "");
      const selectedPrompt = allPrompts.find((p) => p.name === cleanVal);
      if (selectedPrompt && selectedPrompt.loras && selectedPrompt.loras.length > 0) {
        for (const lora of selectedPrompt.loras) {
          const dName = (lora.name || "").replace(/\.safetensors$/, "");
          const str = lora.strength != null ? lora.strength : 1.0;
          const loraW = node.addWidget(
            "text",
            `  🔗 ${dName}`,
            `${str.toFixed(2)}`,
            () => {},
            { serialize: false }
          );
          loraW._pdb_lora_display = true;
          loraW.disabled = true;
        }
      }
    }

    syncSelectionsJSON();
    const computed = node.computeSize?.() || [userWidth, 100];
    node.setSize?.([userWidth, computed[1]]);
    app.graph?.setDirtyCanvas?.(true);
  }

  // Update LoRA display strips when combo selection changes
  async function updateCombinerLoras() {
    // Remove old lora display widgets
    if (node.widgets) {
      node.widgets = node.widgets.filter((w) => !w._pdb_lora_display);
    }

    // Re-add LoRA strips after each combo widget
    const combos = (node.widgets || []).filter((w) => w._pdb_dynamic);
    // We need to insert in order, so collect all widgets to add
    const toInsert = []; // { afterWidget, loraWidgets[] }

    for (const combo of combos) {
      const cleanVal = (combo.value || "").replace(/\s*📖$/, "");
      const cat = combo._pdb_category;
      if (!cleanVal || cleanVal === "(no prompts)" || !cat) continue;

      let prompts = combo._pdb_allPrompts;
      if (!prompts) {
        try {
          prompts = await API.getPrompts(cat);
        } catch (e) {
          continue;
        }
      }

      const prompt = prompts.find((p) => p.name === cleanVal);
      if (!prompt || !prompt.loras || prompt.loras.length === 0) continue;

      toInsert.push({ afterWidget: combo, loras: prompt.loras });
    }

    // Insert LoRA widgets after their combo (in reverse to keep indices stable)
    for (const item of toInsert.reverse()) {
      const idx = node.widgets.indexOf(item.afterWidget);
      if (idx === -1) continue;

      const newWidgets = [];
      for (const lora of item.loras) {
        const dName = (lora.name || "").replace(/\.safetensors$/, "");
        const str = lora.strength != null ? lora.strength : 1.0;
        const w = {
          type: "text",
          name: `  🔗 ${dName}`,
          value: `${str.toFixed(2)}`,
          options: { serialize: false },
          _pdb_lora_display: true,
          disabled: true,
          callback: () => {},
          computeSize: function (width) {
            return [width, 20];
          },
        };
        newWidgets.push(w);
      }
      // Insert after the combo widget
      node.widgets.splice(idx + 1, 0, ...newWidgets);
    }

    const computed = node.computeSize?.() || [userWidth, 100];
    node.setSize?.([userWidth, computed[1]]);
    app.graph?.setDirtyCanvas?.(true);
  }

  function scheduleRebuild() {
    if (rebuildTimer) clearTimeout(rebuildTimer);
    rebuildTimer = setTimeout(async () => {
      if (rebuildPending) return;
      rebuildPending = true;
      try {
        await fullRebuild();
      } finally {
        rebuildPending = false;
      }
    }, 150);
  }

  // Refresh button
  node.addWidget("button", "🔄 Refresh", null, () => scheduleRebuild());

  const origOnConnectionsChange = node.onConnectionsChange;
  node.onConnectionsChange = function (
    type,
    slotIndex,
    isConnected,
    link,
    ioSlot
  ) {
    if (origOnConnectionsChange) origOnConnectionsChange.apply(this, arguments);
    if (type === 1) scheduleRebuild();
  };

  // Periodically check for category changes
  node._pdb_refreshInterval = setInterval(() => {
    if (!node.graph) {
      clearInterval(node._pdb_refreshInterval);
      return;
    }
    let needsRebuild = false;
    for (let i = 0; i < (node.inputs?.length || 0); i++) {
      const inp = node.inputs[i];
      if (!inp.name.startsWith("selection_") || inp.link == null) continue;
      const category = getSourceCategory(i);
      const existing = node.widgets?.find(
        (w) => w._pdb_dynamic && w._pdb_inputName === inp.name
      );
      if (category && (!existing || existing._pdb_category !== category)) {
        needsRebuild = true;
        break;
      }
    }
    if (needsRebuild) scheduleRebuild();
  }, 1000);

  node.setSize?.([userWidth, node.size?.[1] || 100]);
  scheduleRebuild();
}

// ─── Setup ApplyLoraStack ───────────────────────────────────────────────────
const APPLY_LORA_WIDTH = 450;

function setupApplyLoraNode(node) {
  if (node._pdb_apply_setup) return;
  node._pdb_apply_setup = true;
  console.log("[PromptDB] Setting up ApplyLoraStack node", node.id);

  // Find and hide the lora_config text widget
  const configWidget = node.widgets?.find((w) => w.name === "lora_config");
  if (configWidget) {
    configWidget.type = "converted-widget";
    configWidget.computeSize = () => [0, -4];
    configWidget.serializeValue = () => configWidget.value;
  }

  let rebuildPending = false;
  let rebuildTimer = null;

  // Get current config overrides
  function getConfig() {
    try {
      return JSON.parse(configWidget?.value || "{}");
    } catch (e) {
      return {};
    }
  }

  // Save config back to hidden widget
  function saveConfig(config) {
    if (configWidget) configWidget.value = JSON.stringify(config);
  }

  // Trace back through lora_stack input to discover LoRAs
  async function discoverLoras() {
    const loraStackInput = node.inputs?.find((inp) => inp.name === "lora_stack");
    if (!loraStackInput || loraStackInput.link == null) return [];

    const linkInfo = app.graph.links[loraStackInput.link];
    if (!linkInfo) return [];

    const sourceNode = app.graph.getNodeById(linkInfo.origin_id);
    if (!sourceNode) return [];

    const sourceClass = getNodeClass(sourceNode);

    if (sourceClass === "PromptSelector") {
      return await getLorasFromSelector(sourceNode);
    } else if (sourceClass === "PromptCombiner") {
      return await getLorasFromCombiner(sourceNode);
    }

    return [];
  }

  // Get LoRAs from a PromptSelector node
  async function getLorasFromSelector(selectorNode) {
    const catWidget = selectorNode.widgets?.find((w) => w.name === "category");
    const promptWidget = selectorNode.widgets?.find(
      (w) => w.name === "prompt_name"
    );
    if (!catWidget || !promptWidget) return [];

    const cat = catWidget.value;
    const rawName = promptWidget.value || "";
    const cleanName = rawName.replace(/\s*📖$/, "");

    try {
      const prompts = await API.getPrompts(cat);
      const prompt = prompts.find((p) => p.name === cleanName);
      if (prompt && prompt.loras && prompt.loras.length > 0) {
        return prompt.loras.map((l) => ({
          name: l.name || "",
          strength: l.strength != null ? l.strength : 1.0,
        }));
      }
    } catch (e) {
      console.error("[PromptDB] Failed to get LoRAs from selector:", e);
    }
    return [];
  }

  // Get LoRAs from a PromptCombiner node
  async function getLorasFromCombiner(combinerNode) {
    const allLoras = [];
    const seen = new Set();

    // Read all dynamic combo widgets to find selected prompts
    const combos = (combinerNode.widgets || []).filter((w) => w._pdb_dynamic);

    for (const combo of combos) {
      const category = combo._pdb_category;
      const rawVal = combo.value || "";
      const cleanName = rawVal.replace(/\s*📖$/, "");

      if (!category || !cleanName || cleanName === "(no prompts)") continue;
      // Skip muted categories
      if (combo._pdb_muted) continue;

      try {
        const prompts = await API.getPrompts(category);
        const prompt = prompts.find((p) => p.name === cleanName);
        if (prompt && prompt.loras) {
          for (const l of prompt.loras) {
            const name = l.name || "";
            if (name && !seen.has(name)) {
              seen.add(name);
              allLoras.push({
                name,
                strength: l.strength != null ? l.strength : 1.0,
              });
            }
          }
        }
      } catch (e) {
        console.error("[PromptDB] Failed to get LoRAs from combiner:", e);
      }
    }

    return allLoras;
  }

  // Full rebuild of LoRA display widgets
  async function fullRebuild() {
    console.log("[PromptDB] ApplyLoraStack rebuild", node.id);

    // Remove all dynamic LoRA widgets
    if (node.widgets) {
      node.widgets = node.widgets.filter((w) => !w._pdb_lora_widget);
    }

    const loras = await discoverLoras();
    if (loras.length === 0) {
      const emptyW = node.addWidget(
        "text",
        "Status",
        "No LoRAs connected",
        () => {},
        { serialize: false }
      );
      emptyW._pdb_lora_widget = true;
      emptyW.disabled = true;

      const computed = node.computeSize?.() || [APPLY_LORA_WIDTH, 100];
      node.setSize?.([Math.max(computed[0], APPLY_LORA_WIDTH), computed[1]]);
      app.graph?.setDirtyCanvas?.(true);
      return;
    }

    const config = getConfig();

    for (const lora of loras) {
      const displayName = (lora.name || "").replace(/\.safetensors$/, "");
      const override = config[lora.name] || {};
      const isEnabled = override.enabled !== false;
      const str =
        override.strength != null ? override.strength : lora.strength;

      // ── Custom single-line widget: ● name  ◀ 1.00 ▶ ──
      const loraFullName = lora.name; // closure capture
      const widget = {
        type: "pdb_lora",
        name: displayName,
        value: str,
        _pdb_lora_widget: true,
        _pdb_lora_name: loraFullName,
        _pdb_enabled: isEnabled,
        options: { serialize: false },
        last_y: 0,
        computeSize: function () {
          return [node.size[0], 20];
        },
        // ── Custom draw ──
        draw: function (ctx, nodeRef, widgetWidth, posY, H) {
          this.last_y = posY;
          const margin = 15;
          const innerW = widgetWidth - margin * 2;

          // Background bar
          ctx.fillStyle = "#232333";
          ctx.beginPath();
          ctx.roundRect
            ? ctx.roundRect(margin, posY, innerW, H, H * 0.3)
            : ctx.rect(margin, posY, innerW, H);
          ctx.fill();

          // Toggle circle
          const toggleX = margin + 14;
          const toggleY = posY + H * 0.5;
          const r = 5;
          ctx.beginPath();
          ctx.arc(toggleX, toggleY, r, 0, Math.PI * 2);
          ctx.fillStyle = this._pdb_enabled ? "#7cacf8" : "#555";
          ctx.fill();
          if (this._pdb_enabled) {
            ctx.beginPath();
            ctx.arc(toggleX - 1, toggleY - 1, r * 0.35, 0, Math.PI * 2);
            ctx.fillStyle = "rgba(255,255,255,0.25)";
            ctx.fill();
          }

          // LoRA name
          ctx.fillStyle = this._pdb_enabled ? "#ddd" : "#777";
          ctx.font = "12px sans-serif";
          ctx.textAlign = "left";
          ctx.textBaseline = "middle";
          ctx.fillText(this.name, margin + 28, posY + H * 0.5);

          // ◀ value ▶ on right side
          const rightEdge = margin + innerW - 8;
          const valStr = this.value.toFixed(2);

          // ▶ button
          ctx.fillStyle = "#aaa";
          ctx.textAlign = "center";
          ctx.font = "10px sans-serif";
          ctx.fillText("▶", rightEdge, posY + H * 0.5);

          // Value
          ctx.fillStyle = this._pdb_enabled ? "#ddd" : "#777";
          ctx.font = "12px sans-serif";
          ctx.fillText(valStr, rightEdge - 28, posY + H * 0.5);

          // ◀ button
          ctx.fillStyle = "#aaa";
          ctx.font = "10px sans-serif";
          ctx.fillText("◀", rightEdge - 56, posY + H * 0.5);

          ctx.textBaseline = "alphabetic";
        },
        // ── Custom mouse ──
        mouse: function (event, pos, nodeRef) {
          if (
            event.type !== "pointerdown" &&
            event.type !== "mousedown"
          )
            return false;

          const x = pos[0];
          const widgetWidth = nodeRef.size[0];
          const margin = 15;
          const innerW = widgetWidth - margin * 2;
          const rightEdge = margin + innerW - 8;

          // Toggle area (left 50px)
          if (x < margin + 45) {
            this._pdb_enabled = !this._pdb_enabled;
            const cfg = getConfig();
            if (!cfg[this._pdb_lora_name]) cfg[this._pdb_lora_name] = {};
            cfg[this._pdb_lora_name].enabled = this._pdb_enabled;
            saveConfig(cfg);
            app.graph?.setDirtyCanvas?.(true);
            return true;
          }

          // ◀ button (decrease strength)
          if (x > rightEdge - 70 && x < rightEdge - 42) {
            this.value = Math.max(-2, +(this.value - 0.05).toFixed(2));
            const cfg = getConfig();
            if (!cfg[this._pdb_lora_name]) cfg[this._pdb_lora_name] = {};
            cfg[this._pdb_lora_name].strength = this.value;
            saveConfig(cfg);
            app.graph?.setDirtyCanvas?.(true);
            return true;
          }

          // ▶ button (increase strength)
          if (x > rightEdge - 14 && x < rightEdge + 14) {
            this.value = Math.min(2, +(this.value + 0.05).toFixed(2));
            const cfg = getConfig();
            if (!cfg[this._pdb_lora_name]) cfg[this._pdb_lora_name] = {};
            cfg[this._pdb_lora_name].strength = this.value;
            saveConfig(cfg);
            app.graph?.setDirtyCanvas?.(true);
            return true;
          }

          return false;
        },
      };

      node.widgets = node.widgets || [];
      node.widgets.push(widget);
    }

    // ── Add Reset Values button ──
    const resetBtn = node.addWidget(
      "button",
      "🔄 Reset Values",
      null,
      async () => {
        // Clear all overrides
        if (configWidget) configWidget.value = "{}";
        await fullRebuild();
      },
      { serialize: false }
    );
    resetBtn._pdb_lora_widget = true;

    const computed = node.computeSize?.() || [APPLY_LORA_WIDTH, 100];
    node.setSize?.([Math.max(computed[0], APPLY_LORA_WIDTH), computed[1]]);
    app.graph?.setDirtyCanvas?.(true);
  }

  function scheduleRebuild() {
    if (rebuildTimer) clearTimeout(rebuildTimer);
    rebuildTimer = setTimeout(async () => {
      if (rebuildPending) return;
      rebuildPending = true;
      try {
        await fullRebuild();
      } finally {
        rebuildPending = false;
      }
    }, 200);
  }

  // ── Refresh button ──
  node.addWidget("button", "🔄 Refresh", null, () => scheduleRebuild());

  // ── React to connection changes ──
  const origOnConnectionsChange = node.onConnectionsChange;
  node.onConnectionsChange = function (
    type,
    slotIndex,
    isConnected,
    link,
    ioSlot
  ) {
    if (origOnConnectionsChange) origOnConnectionsChange.apply(this, arguments);
    scheduleRebuild();
  };

  // Periodically check for changes in connected source nodes
  let lastLoraFingerprint = "";
  node._pdb_refreshInterval = setInterval(async () => {
    if (!node.graph) {
      clearInterval(node._pdb_refreshInterval);
      return;
    }

    try {
      const loras = await discoverLoras();
      const fingerprint = JSON.stringify(loras.map((l) => l.name).sort());
      if (fingerprint !== lastLoraFingerprint) {
        lastLoraFingerprint = fingerprint;
        scheduleRebuild();
      }
    } catch (e) {}
  }, 1500);

  // Initial setup
  node.setSize?.([APPLY_LORA_WIDTH, node.size?.[1] || 100]);
  scheduleRebuild();
}

// ─── Register extension ─────────────────────────────────────────────────────
app.registerExtension({
  name: "comfy.promptdb",

  async beforeRegisterNodeDef(nodeType, nodeData, _app) {
    if (nodeData.name === "PromptDBManager") {
      const orig = nodeType.prototype.onNodeCreated;
      nodeType.prototype.onNodeCreated = function () {
        orig?.apply(this, arguments);
        setTimeout(() => setupManagerNode(this), 100);
      };
    }
    if (nodeData.name === "PromptSelector") {
      const orig = nodeType.prototype.onNodeCreated;
      nodeType.prototype.onNodeCreated = function () {
        orig?.apply(this, arguments);
        setTimeout(() => setupSelectorNode(this), 100);
      };
    }
    if (nodeData.name === "PromptCombiner") {
      const orig = nodeType.prototype.onNodeCreated;
      nodeType.prototype.onNodeCreated = function () {
        orig?.apply(this, arguments);
        setTimeout(() => setupCombinerNode(this), 100);
      };
    }
    if (nodeData.name === "ApplyLoraStack") {
      const orig = nodeType.prototype.onNodeCreated;
      nodeType.prototype.onNodeCreated = function () {
        orig?.apply(this, arguments);
        setTimeout(() => setupApplyLoraNode(this), 100);
      };
    }
  },

  nodeCreated(node) {
    const cls = getNodeClass(node);
    if (cls === "PromptDBManager")
      setTimeout(() => setupManagerNode(node), 150);
    if (cls === "PromptSelector")
      setTimeout(() => setupSelectorNode(node), 150);
    if (cls === "PromptCombiner")
      setTimeout(() => setupCombinerNode(node), 150);
    if (cls === "ApplyLoraStack")
      setTimeout(() => setupApplyLoraNode(node), 150);
  },
});

console.log("[PromptDB] Extension registered.");