// custom_nodes/YourPkg/js/vibevoice_wrapper_ui.js
import { app } from "../../scripts/app.js";

app.registerExtension({
  name: "vibevoice.wrapper.ui",

  async beforeRegisterNodeDef(nodeType, nodeData) {
    const isWrapper =
      nodeType?.comfyClass === "VibeVoiceTTS_Wrapper" ||
      nodeData?.name === "VibeVoice TTS (Chunked Wrapper)";
    if (!isWrapper) return;

    const origOnCreated = nodeType.prototype.onNodeCreated;
    nodeType.prototype.onNodeCreated = function () {
      origOnCreated?.apply(this, arguments);
      // only set up handlers here; do NOT mutate slots yet
      wireUpHandlers(this);
    };

    function wireUpHandlers(node) {
      const findW = (n) => node.widgets?.find((w) => w.name === n);
      const wNum = findW("num_speakers");
      const wChunk = findW("chunk_lines");
      const wLines = findW("lines_per_chunk");

      function ensureSpeakerInputs(count) {
        // add missing inputs
        for (let i = 1; i <= count; i++) {
          const name = `speaker_${i}_voice`;
          if (node.findInputSlot(name) === -1) node.addInput(name, "AUDIO");
        }
        // remove extras
        for (let i = count + 1; i <= 4; i++) {
          const name = `speaker_${i}_voice`;
          const idx = node.findInputSlot(name);
          if (idx !== -1) node.removeInput(idx);
        }
      }

      // guard: only mutate once node.graph exists (prevents NullGraphError)
      function safeMutate(fn) {
        const doIt = () => {
          if (!node.graph) {
            // defer until the node is actually attached to a graph
            setTimeout(doIt, 0);
            return;
          }
          fn();
          app.graph.setDirtyCanvas(true, true);
        };
        doIt();
      }

      function refresh() {
        const n = Math.max(1, Math.min(4, Number(wNum?.value ?? 1)));
        safeMutate(() => ensureSpeakerInputs(n));

        if (wLines) wLines.hidden = !(wChunk?.value);
      }

      // robust wiring (some frontends only call one of these)
      if (wNum) { wNum.callback = refresh; wNum.onChange = refresh; }
      if (wChunk) { wChunk.callback = refresh; wChunk.onChange = refresh; }

      // don’t call refresh yet; node may not be in graph during configure
      node.__vv_refresh = refresh;
    }
  },

  // Called for brand-new nodes added from the menu (node has a graph here)
  async nodeCreated(node) {
    if (
      node?.comfyClass === "VibeVoiceTTS_Wrapper" ||
      node?.title === "VibeVoice TTS (Chunked Wrapper)"
    ) {
      // next tick to ensure widgets fully exist
      setTimeout(() => node.__vv_refresh?.(), 0);
    }
  },

  // Called when nodes are created as part of loading a workflow
  loadedGraphNode(node) {
    if (
      node?.comfyClass === "VibeVoiceTTS_Wrapper" ||
      node?.title === "VibeVoice TTS (Chunked Wrapper)"
    ) {
      setTimeout(() => node.__vv_refresh?.(), 0);
    }
  },

  // After the graph finishes configuring (safe point to mutate slots)
  async afterConfigureGraph() {
    // final pass in case anything was deferred
    for (const node of app.graph._nodes) {
      if (
        node?.comfyClass === "VibeVoiceTTS_Wrapper" ||
        node?.title === "VibeVoice TTS (Chunked Wrapper)"
      ) {
        node.__vv_refresh?.();
      }
    }
  },

  async setup() {
    console.log("[vibevoice.wrapper.ui] setup complete");
  },
});