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

/**
 * RM Save Lineage - self-growing image inputs.
 *
 * Connect something to the last image slot and a fresh empty one appears below it, so the
 * chain can be extended without ever pressing a button. Its sibling RMMakeImageBatch uses
 * an explicit "Add Image Input" button; here the slot order IS the chain order, so you are
 * always building a sequence and the next slot should already be waiting for you.
 *
 * Middle gaps are deliberately left alone. Disconnecting image_2 of four leaves an empty
 * slot in place rather than renumbering, because the Python side sorts by slot number and
 * skips unwired slots - so 1, 3, 4 is still the chain 1 -> 3 -> 4. Compacting would silently
 * re-point every parent link in the lineage record.
 *
 * Only the tail is managed: there is always exactly one trailing empty slot.
 */

const NODE_NAME = "RMSaveLineage";
const PREFIX = "image_";

/** Image input slots with their parsed number, ordered by that number. */
function imageSlots(node) {
  const slots = [];
  (node.inputs ?? []).forEach((input, index) => {
    if (!input?.name?.startsWith(PREFIX)) return;
    const n = Number(input.name.slice(PREFIX.length));
    if (Number.isInteger(n) && n > 0) slots.push({ index, n, input });
  });
  slots.sort((a, b) => a.n - b.n);
  return slots;
}

const isConnected = (slot) => slot.input.link != null;

/**
 * Guarantee exactly one trailing empty image slot.
 *
 * Idempotent on purpose: LiteGraph fires onConnectionsChange once per restored link while
 * loading a workflow, so this runs many times in a row and must converge, not accumulate.
 */
function ensureTrailingSlot(node) {
  let slots = imageSlots(node);

  // Nothing at all (shouldn't happen - Python declares image_1 - but a hand-edited
  // workflow can arrive without it).
  if (slots.length === 0) {
    node.addInput(`${PREFIX}1`, "IMAGE");
    resize(node);
    return;
  }

  // Trailing run of unconnected slots, counted from the end.
  let trailing = 0;
  for (let i = slots.length - 1; i >= 0; i--) {
    if (isConnected(slots[i])) break;
    trailing++;
  }

  if (trailing === 0) {
    const next = slots[slots.length - 1].n + 1;
    node.addInput(`${PREFIX}${next}`, "IMAGE");
    resize(node);
    return;
  }

  if (trailing > 1) {
    // Remove extras from the end so earlier indices stay valid as we splice.
    // Never drop image_1: the node should always offer at least one slot.
    const removable = slots.slice(slots.length - trailing + 1).filter((s) => s.n !== 1);
    for (let i = removable.length - 1; i >= 0; i--) {
      const idx = node.inputs.findIndex((inp) => inp.name === removable[i].input.name);
      if (idx > -1) node.removeInput(idx);
    }
    resize(node);
  }
}

function resize(node) {
  const size = node.computeSize();
  node.setSize([Math.max(node.size[0], size[0]), size[1]]);
  node.setDirtyCanvas(true, true);
}

app.registerExtension({
  name: "RMAutomation.SaveLineage",

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

    const onNodeCreated = nodeType.prototype.onNodeCreated;
    nodeType.prototype.onNodeCreated = function () {
      const r = onNodeCreated?.apply(this, arguments);
      // Defer: inputs declared by the node definition are not attached yet at this point.
      requestAnimationFrame(() => ensureTrailingSlot(this));
      return r;
    };

    const onConfigure = nodeType.prototype.onConfigure;
    nodeType.prototype.onConfigure = function () {
      const r = onConfigure?.apply(this, arguments);
      // After a workflow loads, links are attached but our slot bookkeeping has not run.
      requestAnimationFrame(() => ensureTrailingSlot(this));
      return r;
    };

    const onConnectionsChange = nodeType.prototype.onConnectionsChange;
    nodeType.prototype.onConnectionsChange = function (type, slotIndex, connected, linkInfo, ioSlot) {
      const r = onConnectionsChange?.apply(this, arguments);

      // LiteGraph.INPUT === 1. Compared numerically because the constant is not always
      // in scope in every frontend build.
      if (type !== 1) return r;
      if (!ioSlot?.name?.startsWith(PREFIX) && !this.inputs?.[slotIndex]?.name?.startsWith(PREFIX)) {
        return r;
      }

      // Deferred so LiteGraph has finished mutating this.inputs before we read it -
      // on disconnect the link is still present during the callback.
      requestAnimationFrame(() => ensureTrailingSlot(this));
      return r;
    };
  },
});