// Audio Wave + Segments — mirrored waveform + playback, a bold 721-frame group grid,
// optional fine user splits inside groups, per-segment notes auto-filled into the `notes`
// box, playhead + click-to-seek. Open the console for [audiowave] logs on trouble.
import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";

function getW(node, name) { return node.widgets?.find((w) => w.name === name); }

function computePeaks(buf, n) {
  const data = buf.getChannelData(0);
  const block = Math.max(1, Math.floor(data.length / n));
  const peaks = new Float32Array(n);
  let max = 1e-6;
  for (let i = 0; i < n; i++) {
    let m = 0; const s = i * block;
    for (let j = 0; j < block && s + j < data.length; j++) { const v = Math.abs(data[s + j]); if (v > m) m = v; }
    peaks[i] = m; if (m > max) max = m;
  }
  for (let i = 0; i < n; i++) peaks[i] /= max;
  return peaks;
}

// Keep a ComfyUI DOM widget at full node width (fixes collapse-to-half on select).
function keepDomWidgetFullWidth(node, container) {
  const GRID_SEL = '[data-testid="node-widgets"], .lg-node-widgets';
  const ROW_SEL = '[data-testid="node-widget"], .lg-node-widget';
  const MAX_MARGIN = 40;
  let enforcing = false, marginLogical = Infinity, gridObserver = null;
  function refFromDom() {
    const grid = container.closest(GRID_SEL); if (!grid) return 0;
    let w = 0;
    for (const row of Array.from(grid.querySelectorAll(ROW_SEL))) { if (row.contains(container)) continue; const c = row.lastElementChild; if (c && c.clientWidth > w) w = c.clientWidth; }
    if (!w && grid.clientWidth > 0) { const dot = grid.querySelector(`${ROW_SEL.split(",")[0]} > :first-child`); w = grid.clientWidth - (dot?.offsetWidth ?? 0); }
    return w;
  }
  function refFromNodeSize(cw) {
    const nodeW = node.size?.[0] ?? 0; if (nodeW <= 0) return 0;
    if (cw > 0) { const m = nodeW - cw; if (m >= 0 && m < marginLogical) marginLogical = Math.min(m, MAX_MARGIN); }
    return Math.round(nodeW - (Number.isFinite(marginLogical) ? marginLogical : MAX_MARGIN / 2));
  }
  function reference(cw) {
    let w = refFromDom(); if (!w) w = refFromNodeSize(cw);
    if (!w) for (const wd of node.widgets ?? []) { const el = wd.inputEl || wd.element; if (el && el !== container && el.offsetWidth > w) w = el.offsetWidth; }
    if (!w && container.parentElement) w = container.parentElement.clientWidth;
    return w;
  }
  function enforce() {
    if (enforcing) return;
    if (!gridObserver) { const grid = container.closest(GRID_SEL); if (grid) { gridObserver = new ResizeObserver(enforce); gridObserver.observe(grid); } }
    const cw = container.clientWidth, ref = reference(cw);
    if (ref > 0 && Math.abs(cw - ref) > 2) { enforcing = true; container.style.width = ref + "px"; requestAnimationFrame(() => { enforcing = false; }); }
  }
  const ro = new ResizeObserver(enforce); ro.observe(container);
  const origOnResize = node.onResize; node.onResize = function (s) { origOnResize?.call(this, s); enforce(); };
  return () => { ro.disconnect(); gridObserver?.disconnect(); };
}

function setupWave(node) {
  const st = { duration: 0, peaks: null, splits: [], audio: new Audio(), playing: false };
  node._wave = st;

  const wrap = document.createElement("div");
  wrap.style.cssText = "display:flex;flex-direction:column;gap:4px;width:100%;";
  const bar = document.createElement("div");
  bar.style.cssText = "display:flex;gap:6px;align-items:center;font-size:10px;color:#bbb;flex-wrap:wrap;";
  const mk = (t) => { const b = document.createElement("button"); b.textContent = t; b.style.cssText = "font-size:10px;padding:1px 6px;"; return b; };
  const playBtn = mk("▶ play"), upBtn = mk("upload"), grpBtn = mk("split @ groups"), clearBtn = mk("clear splits");
  const readout = document.createElement("span"); readout.textContent = "0.00 / 0.00s";
  const hint = document.createElement("span"); hint.textContent = "click=seek · dblclick=add split (snaps to group line/frame) · shift-click a split=remove";
  bar.append(playBtn, upBtn, grpBtn, clearBtn, readout, hint);
  const canvas = document.createElement("canvas");
  canvas.width = 900; canvas.height = 170;
  canvas.style.cssText = "width:100%;height:170px;background:#141418;border-radius:4px;cursor:pointer;";
  wrap.append(bar, canvas);
  node.addDOMWidget("wave", "wave", wrap, { serialize: false, getMinHeight: () => 225 });
  const cleanupWidth = keepDomWidgetFullWidth(node, wrap);
  const origRemoved = node.onRemoved;
  node.onRemoved = function () { try { cleanupWidth(); } catch (e) { /* ignore */ } return origRemoved?.apply(this, arguments); };
  const ctx = canvas.getContext("2d");

  const sjw = getW(node, "segments_json"); if (sjw) { sjw.hidden = true; sjw.computeSize = () => [0, -4]; }

  const fps = () => Math.max(1, Number(getW(node, "fps")?.value) || 24);
  const groupFrames = () => Math.max(1, Number(getW(node, "group_frames")?.value) || 721);  // guard: never 0
  const t2x = (t) => (st.duration ? (t / st.duration) * canvas.width : 0);
  const x2t = (x) => (st.duration ? (x / canvas.width) * st.duration : 0);

  function groupBounds() {
    const step = groupFrames() / fps(); const arr = [];
    if (!st.duration || step <= 0) return arr;
    for (let t = step; t < st.duration - 1e-6; t += step) arr.push(t);
    return arr;
  }
  function segPoints() {                               // [0] + user splits (NOT the group lines)
    const set = new Set([0]);
    for (const s of st.splits) if (s > 0 && s < st.duration) set.add(Math.round(s * 100) / 100);
    return Array.from(set).sort((a, b) => a - b);
  }
  const nSegs = () => Math.max(0, segPoints().length - 1);  // N splits -> N segments (tail excluded)

  function serialize() {                               // write ONLY user splits (small + stable)
    const w = getW(node, "segments_json"); if (!w) return;
    w.value = JSON.stringify(st.splits.map((s) => ({ start_s: Math.round(s * 100) / 100 })));
    w.callback?.(w.value);
  }
  function parseNotes(txt) {
    const seg = {}, glob = [];
    (txt || "").split("\n").forEach((l) => { const m = l.match(/^\s*(?:seg(?:ment)?|s)?\s*(\d+)\s*[:)]\s*(.*)$/i); if (m) seg[parseInt(m[1])] = m[2]; else if (l.trim()) glob.push(l); });
    return { seg, glob };
  }
  function syncNotesBox() {                            // one segN: line per segment, preserve notes
    const nw = getW(node, "notes"); if (!nw) return;
    const { seg, glob } = parseNotes(nw.value);
    const n = nSegs(); const lines = [];
    for (let i = 1; i <= n; i++) lines.push(`seg${i}: ${seg[i] !== undefined ? seg[i] : ""}`);
    const val = lines.concat(glob).join("\n");
    if (nw.value !== val) { nw.value = val; nw.callback?.(val); }
  }

  function selRange() {
    const v = (getW(node, "segment_select")?.value || "").trim();
    const m = v.match(/^(\d+)\s*[-:]\s*(\d+)$/); if (m) return [+m[1], +m[2]].sort((a, b) => a - b);
    if (/^\d+$/.test(v) && v !== "0") return [+v, +v];
    return null;
  }

  function draw() {
    ctx.fillStyle = "#141418"; ctx.fillRect(0, 0, canvas.width, canvas.height);
    const H = canvas.height, mid = H / 2;
    const pts = segPoints();                            // [0, split1, split2, ...]; seg k = pts[k-1]..pts[k]
    const sel = selRange();
    if (sel) {                                          // shade selected segment range
      const a = pts[sel[0] - 1], b = sel[1] < pts.length ? pts[sel[1]] : st.duration;
      if (a !== undefined) { ctx.fillStyle = "#26364f"; ctx.fillRect(t2x(a), 0, t2x(b) - t2x(a), H); }
    }
    if (st.peaks) {                                     // mirrored waveform around centre
      ctx.strokeStyle = "#3c8cdc"; ctx.beginPath();
      const n = st.peaks.length;
      for (let i = 0; i < n; i++) { const x = (i / n) * canvas.width, h = st.peaks[i] * (H * 0.44); ctx.moveTo(x, mid - h); ctx.lineTo(x, mid + h); }
      ctx.stroke();
    }
    ctx.strokeStyle = "#ececf2"; ctx.lineWidth = 2;     // bold group grid
    for (const g of groupBounds()) { const x = t2x(g); ctx.beginPath(); ctx.moveTo(x, 0); ctx.lineTo(x, H); ctx.stroke(); }
    ctx.lineWidth = 1;
    const { seg } = parseNotes(getW(node, "notes")?.value);
    ctx.font = "10px monospace";
    const nseg = nSegs();
    let lastLabelX = -999, lastNoteX = -999;            // skip labels that would overlap
    pts.forEach((s, i) => {                              // split line; label seg i+1 (skip tail)
      const x = t2x(s);
      ctx.strokeStyle = "#7ec8a0"; ctx.beginPath(); ctx.moveTo(x, 0); ctx.lineTo(x, H); ctx.stroke();
      if (i >= nseg) return;
      if (x - lastLabelX > 26) { ctx.fillStyle = "#fff"; ctx.fillText(`S${i + 1}`, x + 2, 11); lastLabelX = x; }
      if (seg[i + 1] && x - lastNoteX > 60) { ctx.fillStyle = "#ffd27a"; ctx.fillText(seg[i + 1].slice(0, 16), x + 2, H - 4); lastNoteX = x; }
    });
    ctx.strokeStyle = "#ff5a3c"; const px = t2x(st.audio.currentTime || 0);  // playhead
    ctx.beginPath(); ctx.moveTo(px, 0); ctx.lineTo(px, H); ctx.stroke();
    readout.textContent = `${(st.audio.currentTime || 0).toFixed(2)} / ${st.duration.toFixed(2)}s`;
  }

  async function loadFile(name) {
    if (!name || name.startsWith("(")) return;
    const url = `/view?filename=${encodeURIComponent(name)}&type=input&subfolder=`;
    try {
      const ab = await (await fetch(url)).arrayBuffer();
      const ac = new (window.AudioContext || window.webkitAudioContext)();
      const buf = await ac.decodeAudioData(ab.slice(0));
      st.duration = buf.duration; st.peaks = computePeaks(buf, canvas.width);
      st.audio.src = url; syncNotesBox(); draw();
    } catch (e) { console.error("[audiowave] could not load/decode", name, e); }
  }

  const tol = () => x2t(6);
  canvas.addEventListener("mousedown", (e) => {
    const r = canvas.getBoundingClientRect();
    const t = x2t(((e.clientX - r.left) / r.width) * canvas.width);
    if (e.shiftKey) {                                   // shift-click = remove nearest user split
      let bi = -1, bd = 1e9; st.splits.forEach((s, i) => { const dd = Math.abs(s - t); if (dd < tol() && dd < bd) { bd = dd; bi = i; } });
      if (bi >= 0) { st.splits.splice(bi, 1); serialize(); syncNotesBox(); draw(); }
      return;
    }
    st.audio.currentTime = Math.max(0, Math.min(st.duration, t)); draw();  // click = seek
  });
  function snap(t) {                                    // snap to a group line if near, else to a frame
    const gs = groupFrames() / fps();
    const gt = Math.round(t / gs) * gs;
    if (gt > 0.02 && gt < st.duration - 0.02 && Math.abs(t2x(gt) - t2x(t)) < 9) return gt;
    return Math.round(t * fps()) / fps();
  }
  canvas.addEventListener("dblclick", (e) => {          // add a user split (snapped)
    const r = canvas.getBoundingClientRect();
    const t = snap(x2t(((e.clientX - r.left) / r.width) * canvas.width));
    if (t > 0.02 && t < st.duration - 0.02) { st.splits.push(t); serialize(); syncNotesBox(); draw(); }
  });

  const loop = () => { if (!st.playing) return; draw(); requestAnimationFrame(loop); };
  playBtn.onclick = () => {
    if (st.audio.paused) { st.audio.play().catch((e) => console.error("[audiowave] play", e)); st.playing = true; playBtn.textContent = "⏸ pause"; loop(); }
    else { st.audio.pause(); st.playing = false; playBtn.textContent = "▶ play"; }
  };
  st.audio.onended = () => { st.playing = false; playBtn.textContent = "▶ play"; draw(); };
  clearBtn.onclick = () => { st.splits = []; serialize(); syncNotesBox(); draw(); };  // wipe stale user splits
  grpBtn.onclick = () => {                              // one split on every 721 group line
    const set = new Set(st.splits.map((s) => Math.round(s * 100) / 100));
    for (const g of groupBounds()) set.add(Math.round(g * 100) / 100);
    st.splits = Array.from(set).sort((a, b) => a - b);
    serialize(); syncNotesBox(); draw();
  };

  upBtn.onclick = () => {
    const inp = document.createElement("input"); inp.type = "file"; inp.accept = "audio/*";
    inp.onchange = async () => {
      const f = inp.files[0]; if (!f) return;
      const fd = new FormData(); fd.append("image", f, f.name); fd.append("type", "input");
      try {
        const res = await api.fetchApi("/upload/image", { method: "POST", body: fd });
        const j = await res.json(); const name = j.name || f.name;
        const w = getW(node, "audio"); if (w) { if (!w.options.values.includes(name)) w.options.values.push(name); w.value = name; }
        await loadFile(name);
      } catch (e) { console.error("[audiowave] upload failed (drop the file in ComfyUI/input instead)", e); }
    };
    inp.click();
  };

  const hook = (name, fn) => { const w = getW(node, name); if (w) { const cb = w.callback; w.callback = function () { const r = cb ? cb.apply(this, arguments) : undefined; fn(); return r; }; } };
  hook("audio", () => loadFile(getW(node, "audio")?.value));
  hook("fps", () => { syncNotesBox(); draw(); });
  hook("group_frames", () => { syncNotesBox(); draw(); });
  hook("segment_select", () => draw());
  try { const sj = getW(node, "segments_json"); if (sj?.value) st.splits = (JSON.parse(sj.value) || []).map((s) => s.start_s).filter((s) => s > 0); } catch (e) { /* ignore */ }
  const aw = getW(node, "audio"); if (aw?.value) loadFile(aw.value);
  draw();
}

app.registerExtension({
  name: "promptcalib.audiowave",
  async beforeRegisterNodeDef(nodeType, nodeData) {
    if (nodeData?.name !== "AudioWaveSegments") return;
    const orig = nodeType.prototype.onNodeCreated;
    nodeType.prototype.onNodeCreated = function () {
      const r = orig?.apply(this, arguments);
      try { setupWave(this); } catch (e) { console.error("[audiowave] setup failed", e); }
      return r;
    };
  },
});