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

const NODE_NAME = "DumasH3PromptCurator";
const EXPECTED_OUTPUTS = [
  "prompt",
  "ref_image_1",
  "ref_image_2",
  "ref_image_3",
  "ref_image_4",
  "ref_image_5",
  "ref_image_6",
  "ref_image_7",
  "ref_image_8",
  "ref_image_9",
  "reference_count",
  "debug",
  "anchor",
  "sounds",
  "bgm",
  "original_ref_1",
  "original_ref_2",
  "original_ref_3",
  "original_ref_4",
  "original_ref_5",
  "original_ref_6",
  "original_ref_7",
  "original_ref_8",
  "original_ref_9",
  "compiled_ref_description_1",
  "compiled_ref_description_2",
  "compiled_ref_description_3",
  "compiled_ref_description_4",
  "compiled_ref_description_5",
  "compiled_ref_description_6",
  "compiled_ref_description_7",
  "compiled_ref_description_8",
  "compiled_ref_description_9",
];
const EXPECTED_NAMES = new Set(EXPECTED_OUTPUTS);

function pruneStaleOutputs(node) {
  if (!Array.isArray(node.outputs)) return;

  const byName = new Map();
  for (const output of node.outputs) {
    if (!output?.name || !EXPECTED_NAMES.has(output.name) || byName.has(output.name)) continue;
    byName.set(output.name, output);
  }

  const nextOutputs = [];
  for (const name of EXPECTED_OUTPUTS) {
    const existing = byName.get(name);
    if (existing) {
      nextOutputs.push(existing);
    }
  }

  if (nextOutputs.length && nextOutputs.length !== node.outputs.length) {
    node.outputs = nextOutputs;
    node.size = node.computeSize?.() || node.size;
    node.setDirtyCanvas?.(true, true);
  }
}

app.registerExtension({
  name: "Dumas.H3PromptCuratorOutputs",
  async beforeRegisterNodeDef(nodeType, nodeData) {
    if (nodeData?.name !== NODE_NAME) return;

    const originalOnNodeCreated = nodeType.prototype.onNodeCreated;
    const originalOnConfigure = nodeType.prototype.onConfigure;

    nodeType.prototype.onNodeCreated = function onNodeCreated() {
      const created = originalOnNodeCreated?.apply(this, arguments);
      pruneStaleOutputs(this);
      return created;
    };

    nodeType.prototype.onConfigure = function onConfigure() {
      const configured = originalOnConfigure?.apply(this, arguments);
      pruneStaleOutputs(this);
      return configured;
    };
  },
});