// AI summaries per article section, running fully on-device via WebGPU.
// Model runtime vendored from huggingface.co/spaces/webml-community/lfm2-webgpu-kernels
// (Lfm2Mobile: GGUF loader + custom WGSL kernels + generation loop, all in lfm2_5.js).
import { Lfm2Mobile, DEFAULT_MODEL_ID } from "./lfm2_5.js";

// A "section" is an H1 or H2 plus every direct-child element after it up to the next
// H1/H2 (mirrors the boundaries the .article-sidebar TOC already uses — H3s are
// sub-headings within their parent section, not their own summarizable unit).
const HEADING_SELECTOR = ":scope > h1, :scope > h2";

const MIN_CHARS = 60; // skip empty/near-empty sections (e.g. a heading right before the next one)
const MAX_NEW_TOKENS = 512;
const CIRC = 2 * Math.PI * 13; // matches the r=13 <circle> in renderRing()

// Model + its download/warmup are shared across every section on the page: only the
// first click anywhere pays the load cost. `progressListeners` fans load progress out to
// every ring currently waiting on it (a section clicked mid-download gets live updates,
// not just the one that triggered the load). `genTail` serializes generate() calls, since
// the model can only run one generation at a time.
let model = null;
let loadPromise = null;
const progressListeners = new Set();
let genTail = Promise.resolve();

document.addEventListener("DOMContentLoaded", init);

function init() {
  const article = document.querySelector(".article");
  if (!article) return;

  // Read the section boundaries from the ORIGINAL tree before any wrapping happens below
  // (wrapping a heading only reparents the heading itself, not its following siblings, but
  // computing every section's content list up front keeps this a clean read-then-write pass).
  // Also: must run before script.js wraps <pre> blocks into .code-collapsible — relies on
  // this file's <script type="module"> tag appearing before ../script.js in each post's
  // HTML, which makes it execute first (defer/module scripts run in document order).
  const headings = Array.from(article.querySelectorAll(HEADING_SELECTOR));

  headings
    .map((heading, i) => ({ heading, contentEls: collectSection(heading, headings[i + 1] || null) }))
    .filter(({ contentEls }) => sectionText(contentEls).length >= MIN_CHARS)
    .forEach(({ heading, contentEls }) => setupSection(heading, contentEls));
}

function collectSection(heading, until) {
  const els = [];
  let node = heading.nextElementSibling;
  while (node && node !== until) {
    els.push(node);
    node = node.nextElementSibling;
  }
  return els;
}

function sectionText(contentEls) {
  return contentEls.map(extractElementText).filter(Boolean).join("\n\n");
}

// textContent of a stripped clone so we never touch the live DOM; drops KaTeX's hidden
// MathML annotation clone (raw LaTeX source) so math-heavy sections don't feed the model
// duplicated/garbled text. Tables and lists get a lightweight text layout so cells/items
// don't run together; code blocks keep their original whitespace.
function extractElementText(el) {
  const clone = el.cloneNode(true);
  clone.querySelectorAll(".katex-mathml").forEach((n) => n.remove());

  if (clone.tagName === "TABLE") return extractTableText(clone);
  if (clone.tagName === "UL" || clone.tagName === "OL") return extractListText(clone);
  if (clone.tagName === "PRE") return clone.textContent.trim();
  return clone.textContent.replace(/\s+/g, " ").trim();
}

function extractTableText(table) {
  return Array.from(table.querySelectorAll("tr"))
    .map((tr) =>
      Array.from(tr.querySelectorAll("th,td"))
        .map((cell) => cell.textContent.replace(/\s+/g, " ").trim())
        .join(" | ")
    )
    .join("\n");
}

function extractListText(list) {
  return Array.from(list.querySelectorAll(":scope > li"))
    .map((li) => "- " + li.textContent.replace(/\s+/g, " ").trim())
    .join("\n");
}

function setupSection(heading, contentEls) {
  const wrap = document.createElement("div");
  wrap.className = "ai-sum-wrap";
  heading.parentNode.insertBefore(wrap, heading);
  wrap.appendChild(heading);

  const btn = document.createElement("button");
  btn.type = "button";
  btn.className = "ai-sum-btn";
  btn.setAttribute("aria-label", "Summarize this section");
  btn.title = btn.getAttribute("aria-label");
  btn.innerHTML = SPARK_ICON;
  heading.appendChild(btn);

  let note = null;
  let state = "idle"; // idle | busy | done | error

  btn.addEventListener("click", () => {
    if (state === "busy") return;
    if (state === "done") {
      toggleNote();
      return;
    }
    startSummary(); // idle or error -> (re)start
  });

  function ensureNote() {
    if (note) return note;
    note = document.createElement("div");
    note.className = "ai-sum-note";

    const toggle = document.createElement("button");
    toggle.type = "button";
    toggle.className = "ai-sum-note-toggle";
    toggle.setAttribute("aria-expanded", "true");
    toggle.innerHTML = '<span class="ai-sum-arrow">▼</span><span>AI summary</span>';
    toggle.addEventListener("click", toggleNote);

    const body = document.createElement("div");
    body.className = "ai-sum-note-body";
    body.setAttribute("aria-live", "polite");

    note.appendChild(toggle);
    note.appendChild(body);
    wrap.appendChild(note); // right after the heading, before the section's own content
    return note;
  }

  function toggleNote() {
    if (!note) return;
    const collapsed = note.classList.toggle("collapsed");
    note.querySelector(".ai-sum-note-toggle").setAttribute("aria-expanded", String(!collapsed));
  }

  async function startSummary() {
    state = "busy";
    btn.classList.add("is-busy");
    btn.disabled = true;

    ensureNote();
    note.classList.remove("collapsed");
    note.querySelector(".ai-sum-note-toggle").setAttribute("aria-expanded", "true");
    const body = note.querySelector(".ai-sum-note-body");
    const ring = renderRing(body, "Requesting GPU device…");

    const onProgress = (event) => applyLoadProgress(ring, event);

    try {
      if (!navigator.gpu) {
        throw new Error("Summaries need a WebGPU browser (recent Chrome or Edge).");
      }

      if (!model) {
        progressListeners.add(onProgress);
        try {
          model = await getModel();
        } finally {
          progressListeners.delete(onProgress);
        }
      }

      ring.setLabel("Generating summary…");
      ring.setDetail("");
      ring.setIndeterminate();

      function limitWords(text, maxWords) {
  return text.split(/\s+/).slice(0, maxWords).join(" ");
}

      const input = limitWords(sectionText(contentEls), 130);

      const prompt = "Summarize this section directly:\n\n" + input;

      let textEl = null; // created on the first streamed chunk, once we drop the ring
      const summary = await enqueueGenerate(() =>
        runGenerate(prompt, (partial) => {
          if (!textEl) {
            ring.destroy();
            textEl = renderStreamText(body);
          }
          textEl.textContent = partial;
        })
      );

      const finalEl = textEl || renderStreamText(body); // stream yielded nothing before completing
      finalEl.textContent = summary || "(empty response)";
      finalEl.classList.remove("ai-sum-streaming");
      state = "done";
    } catch (err) {
      console.error(err);
      ring.destroy();
      body.innerHTML = "";
      const errEl = document.createElement("p");
      errEl.className = "ai-sum-error";
      errEl.textContent = "Couldn't generate a summary — " + (err?.message ?? String(err));
      body.appendChild(errEl);
      state = "error";
    } finally {
      btn.classList.remove("is-busy");
      btn.disabled = false;
    }
  }
}

function getModel() {
  if (!loadPromise) {
    loadPromise = (async () => {
      const m = await Lfm2Mobile.load(DEFAULT_MODEL_ID, {
        onProgress: (event) => {
          for (const fn of progressListeners) fn(event);
        },
      });
      for (const fn of progressListeners) fn({ status: "warmup" });
      await m.warmup();
      return m;
    })().catch((err) => {
      loadPromise = null; // let the next click retry the load
      throw err;
    });
  }
  return loadPromise;
}

// Runs one generation exclusively — chained onto the shared tail so two sections clicked
// close together never call generate() on the model concurrently.
function enqueueGenerate(fn) {
  const run = genTail.then(fn, fn);
  genTail = run.then(
    () => {},
    () => {}
  );
  return run;
}

// Streams the reply, but caps how often onChunk actually touches the DOM — small models
// can emit far faster than the screen (or a screen reader on the note's aria-live region)
// can usefully keep up with.
const STREAM_RENDER_MS = 90;

async function runGenerate(prompt, onChunk) {
  const stream = model.generate([{ role: "user", content: prompt }], {
    maxNewTokens: MAX_NEW_TOKENS,
  });
  let full = "";
  let lastRenderAt = 0;
  for await (const chunk of stream) {
    full = chunk.text;
    const now = performance.now();
    if (now - lastRenderAt >= STREAM_RENDER_MS) {
      lastRenderAt = now;
      onChunk?.(full);
    }
  }
  return full.trim();
}

function renderStreamText(container) {
  container.innerHTML = "";
  const el = document.createElement("span");
  el.className = "ai-sum-streaming";
  container.appendChild(el);
  return el;
}

// Maps Lfm2Mobile.load()'s onProgress events onto a 0..0.9 ring (0.9..1 is warmup,
// driven separately below), mirroring the cost-weighted mapping the kernels demo uses:
// the GGUF download dwarfs everything else, so it owns most of the bar.
function applyLoadProgress(ring, event) {
  if (event.status === "warmup") {
    ring.setLabel("Warming up kernels…");
    ring.setDetail("");
    ring.setIndeterminate();
    return;
  }

  const frac = Number.isFinite(event.fraction) ? Math.max(0, Math.min(1, event.fraction)) : null;
  const labels = {
    init: "Requesting GPU device…",
    tokenizer: "Loading tokenizer…",
    weights: event.fromCache ? "Loading cached model…" : "Downloading model…",
    ready: "Preparing GPU weights…",
  };
  ring.setLabel(labels[event.status] ?? String(event.status));

  if (event.status === "weights") {
    if (frac !== null) ring.setProgress(0.08 + frac * 0.82);
    if (Number.isFinite(event.loaded) && Number.isFinite(event.total)) {
      ring.setDetail(formatBytes(event.loaded) + " / " + formatBytes(event.total));
    } else {
      ring.setDetail("");
    }
  } else if (event.status === "tokenizer") {
    ring.setProgress(0.06);
    ring.setDetail("");
  } else if (event.status === "ready") {
    ring.setProgress(0.9);
    ring.setDetail("");
  } else {
    ring.setProgress(0.02);
    ring.setDetail("");
  }
}

function formatBytes(n) {
  if (n < 1024) return n + " B";
  const units = ["KB", "MB", "GB"];
  let u = -1;
  do {
    n /= 1024;
    u++;
  } while (n >= 1024 && u < units.length - 1);
  return n.toFixed(1) + " " + units[u];
}

// A small circular (never spinning) progress indicator: a determinate ring driven by
// setProgress() while we know a real fraction (model download), and an indeterminate
// "breathing" ring — stroke-dashoffset oscillating via CSS keyframes — once we don't
// (warmup, token generation).
function renderRing(container, initialLabel) {
  container.innerHTML = "";
  const wrapEl = document.createElement("div");
  wrapEl.className = "ai-sum-progress";
  wrapEl.innerHTML =
    '<span class="ai-sum-ring">' +
    '<svg viewBox="0 0 32 32">' +
    '<circle class="ai-sum-ring-track" cx="16" cy="16" r="13"/>' +
    '<circle class="ai-sum-ring-fill" cx="16" cy="16" r="13" ' +
    'stroke-dasharray="' +
    CIRC.toFixed(2) +
    '" stroke-dashoffset="' +
    CIRC.toFixed(2) +
    '"/>' +
    "</svg>" +
    "</span>" +
    '<span class="ai-sum-progress-text">' +
    '<span class="ai-sum-progress-label"></span>' +
    '<span class="ai-sum-progress-detail"></span>' +
    "</span>";
  container.appendChild(wrapEl);

  const fill = wrapEl.querySelector(".ai-sum-ring-fill");
  const ringEl = wrapEl.querySelector(".ai-sum-ring");
  const labelEl = wrapEl.querySelector(".ai-sum-progress-label");
  const detailEl = wrapEl.querySelector(".ai-sum-progress-detail");

  const api = {
    setLabel(text) {
      labelEl.textContent = text;
    },
    setDetail(text) {
      detailEl.textContent = text || "";
    },
    setProgress(frac) {
      ringEl.classList.remove("is-indeterminate");
      fill.style.strokeDashoffset = (CIRC * (1 - Math.max(0, Math.min(1, frac)))).toFixed(2);
    },
    setIndeterminate() {
      ringEl.classList.add("is-indeterminate");
    },
    destroy() {
      wrapEl.remove();
    },
  };
  api.setLabel(initialLabel);
  api.setProgress(0.02);
  return api;
}

const SPARK_ICON =
  '<svg viewBox="0 0 24 24" width="14" height="14" fill="currentColor" aria-hidden="true">' +
  '<path d="M12 2.5c.3 3.4 1.1 5.7 2.4 7 1.3 1.3 3.6 2.1 7 2.4-3.4.3-5.7 1.1-7 2.4-1.3 1.3-2.1 3.6-2.4 7' +
  "-.3-3.4-1.1-5.7-2.4-7-1.3-1.3-3.6-2.1-7-2.4 3.4-.3 5.7-1.1 7-2.4 1.3-1.3 2.1-3.6 2.4-7z\"/>" +
  "</svg>";