// Drop-in workaround for ComfyUI_frontend #12443.
//
// WidgetLegacy.vue's draw() unconditionally writes the Vue container's
// clientWidth onto widget.width on every frame. In LiteGraph mode this
// pollutes the widget object and the canvas renderer's
// `widget.width || nodeWidth` fallback then ignores node.size[0],
// so widgets stop tracking node resizes after a move.
//
// We can't edit the compiled Vue code, so we neutralize the *effect*:
// replace widget.width with a getter/setter that drops writes while
// the patch is active and LiteGraph mode is in control.
//
// Patch is installed when a LegacyWidgetWidthFix node is added to the
// graph. New nodes are repaired automatically while it is present, and
// the fix node also exposes a manual repair button as a fallback.

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

const NODE_NAME = "LegacyWidgetWidthFix";
const WIDTH_KEY = "__lwwf_width";
const GUARDED = "__lwwf_guarded";
const REPAIR_BUTTON = "__lwwf_repair_button";

let activeCount = 0;
let hooksInstalled = false;
const origMethods = {};
// addDOMWidget and addCustomWidget bypass addWidget — they push directly
// onto node.widgets. Hook all three so we catch every code path that
// produces a widget, regardless of which custom node pack created it.
const HOOK_METHODS = ["addWidget", "addDOMWidget", "addCustomWidget"];

function inVueMode() {
  // LiteGraph.vueNodesMode is the upstream flag the proper fix gates on.
  // If it's undefined (older builds, tree-shaken, etc.) we assume LG mode
  // and drop the write — that matches the regression we're fixing.
  const LG = globalThis.LiteGraph;
  return !!(LG && LG.vueNodesMode);
}

function guardWidget(w) {
  if (!w || w[GUARDED]) return;
  w[WIDTH_KEY] = w.width;
  Object.defineProperty(w, "width", {
    configurable: true,
    enumerable: true,
    get() {
      return this[WIDTH_KEY];
    },
    set(v) {
      if (!patchActive() || inVueMode()) {
        this[WIDTH_KEY] = v;
      }
      // else: drop the write — this is the bug we're suppressing
    },
  });
  Object.defineProperty(w, GUARDED, {
    value: true,
    enumerable: false,
    configurable: true,
  });
}

function repairWidget(w) {
  if (!w) return;
  guardWidget(w);

  // Clear the stale persisted value so renderers fall back to node.size[0].
  // Must be `undefined`, not 0 — DomWidgets.vue uses `width ?? nodeWidth`
  // (nullish), which would not fall back from 0. Canvas renderer uses
  // `width || nodeWidth` (falsy), which falls back from either. undefined
  // works for both.
  if (!inVueMode()) w[WIDTH_KEY] = undefined;
}

function unguardWidget(w) {
  if (!w || !w[GUARDED]) return;
  const v = w[WIDTH_KEY];
  delete w.width;
  delete w[WIDTH_KEY];
  delete w[GUARDED];
  w.width = v;
}

function patchActive() {
  return activeCount > 0;
}

function hookWidgetMethods() {
  if (hooksInstalled) return;
  const proto = globalThis.LGraphNode?.prototype;
  if (!proto) return;
  for (const name of HOOK_METHODS) {
    const fn = proto[name];
    if (typeof fn !== "function") continue;
    origMethods[name] = fn;
    proto[name] = function patched(...args) {
      const w = fn.apply(this, args);
      if (w) {
        guardWidget(w);
        if (patchActive()) repairWidget(w);
      }
      return w;
    };
  }
  hooksInstalled = true;
}

function forEachWidget(fn) {
  const nodes = app.graph?._nodes ?? [];
  for (const n of nodes) {
    const ws = n.widgets;
    if (!ws) continue;
    for (const w of ws) fn(w, n);
  }
}

function repairNode(node) {
  for (const w of node?.widgets ?? []) repairWidget(w);
}

function repairAll() {
  if (!patchActive()) return;
  forEachWidget((w) => repairWidget(w));
  app.graph?.setDirtyCanvas(true, true);
}

function addRepairButton(node) {
  if (!node || node[REPAIR_BUTTON]) return;
  const button = node.addWidget?.(
    "button",
    "Repair widget widths now",
    null,
    () => repairAll(),
    { serialize: false },
  );
  if (button) node[REPAIR_BUTTON] = button;
}

function isFixNode(node) {
  return node?.comfyClass === NODE_NAME
    || node?.constructor?.comfyClass === NODE_NAME
    || node?.type === NODE_NAME;
}

function install() {
  hookWidgetMethods();
  repairAll();
}

function uninstall() {
  forEachWidget((w) => unguardWidget(w));
  app.graph?.setDirtyCanvas(true, true);
}

app.registerExtension({
  name: "LegacyWidgetWidthFix",

  nodeCreated(node) {
    if (isFixNode(node)) addRepairButton(node);
    if (!patchActive()) return;

    repairNode(node);

    // Other extensions can append widgets in their own nodeCreated hooks.
    // Re-scan once those synchronous hooks and the current render pass finish.
    globalThis.requestAnimationFrame?.(() => {
      if (patchActive() && app.graph?._nodes?.includes(node)) repairNode(node);
    });
  },

  beforeRegisterNodeDef(nodeType, nodeData) {
    if (nodeData?.name !== NODE_NAME) return;

    const origOnAdded = nodeType.prototype.onAdded;
    nodeType.prototype.onAdded = function (...args) {
      origOnAdded?.apply(this, args);
      addRepairButton(this);
      activeCount += 1;
      if (activeCount === 1) install();
    };

    const origOnRemoved = nodeType.prototype.onRemoved;
    nodeType.prototype.onRemoved = function (...args) {
      origOnRemoved?.apply(this, args);
      activeCount = Math.max(0, activeCount - 1);
      if (activeCount === 0) uninstall();
    };
  },
});