vscode-cowriting-plugin/src/previewAnnotations.ts

/**
 * previewAnnotations — Task 7 (D3/D21, PUC-3): a pure, vscode-free markdown-it
 * plugin that annotates the BUILT-IN VS Code Markdown preview with the F10
 * authorship/change-mark vocabulary (`.cw-ins-claude` / `.cw-ins-human` /
 * `.cw-del`) — the native-surfaces migration's replacement for the sunset
 * webview's review body (spec §6.4/§6.10). Reuses the shipped sentinel
 * discipline (`injectSentinels`/`sentinelsToSpans`, #33/#47-hardened) instead
 * of inventing a new one: `annotateSource` computes PUA-sentinel-wrapped
 * markdown SOURCE (never HTML) against the F6/F7 baseline + F3 spans + F4
 * pending proposals; `cowritingMarkdownItPlugin` wires that transform into the
 * shared `markdown-it` instance (the SAME instance VS Code's built-in preview
 * uses via `markdown.markdownItPlugins`) as (a) a core rule — swap `state.src`
 * before `normalize` — and (b) a full-render wrapper that turns surviving
 * sentinel pairs into `<span class="cw-…">`.
 *
 * Implementation note (resolved drift from the migration plan's Task 7 sketch,
 * which described (b) as "a text renderer rule"): `sentinelsToSpans` is a
 * stateful, single-pass walk over one HTML string — installing it as a
 * PER-TOKEN `renderer.rules.text` hook would reset that state at every text
 * token boundary and silently break the "survives emphasis/tag boundaries"
 * invariant (#33 CASE3) for any span whose sentinel-open and sentinel-close
 * land in DIFFERENT text tokens (e.g. either side of a `**bold**` run — the
 * exact scenario CASE3 exists to cover). Wrapping `md.renderer.render` instead
 * runs the walk ONCE over the fully-assembled HTML, exactly like the F9/F10
 * `colorByAuthor` precedent (`render(injected)` — a whole-block HTML string,
 * not a token fragment) — this is what actually delivers the stated invariant,
 * robust to callers that invoke `md.render()` directly or `md.parse()` +
 * `md.renderer.render()` separately (VS Code's own scroll-sync `data-line`
 * injection uses the latter).
 *
 * `host.inputsFor(env)` is the ONE thin, vscode-touching seam (owned by
 * `extension.ts`): `env.currentDocument` is an OBSERVED VS Code preview-engine
 * behavior, not a documented contract, so the host falls back to the last
 * actively-coedited document when it is absent (see extension.ts's
 * `lastCoeditedUri`).
 *
 * Deliberate choice: the mermaid-fence rendering (the `options.highlight`
 * override installed below) applies to EVERY markdown preview, coedited or
 * not — it is not gated on coediting state. Gating it would make a diagram's
 * rendering flip (mermaid <-> plain fence) the instant a doc enters or exits
 * coediting, which is a worse surprise than always rendering mermaid diagrams
 * as diagrams; only the change-tracking OVERLAY on top (`buildMermaidQueue`'s
 * diff augmentation) is gated on having a baseline to diff against.
 */
import type MarkdownIt from "markdown-it";
import { diffMermaid } from "./mermaidDiff";
import {
  diffBlocks,
  injectSentinels,
  landedTextOf,
  MERMAID_LEGEND,
  mermaidFenceBody,
  sentinelsToSpans,
  splitBlocksWithRanges,
  wordEditHunks,
  type AuthorSpan,
  type ProposalView,
  type TaggedSpan,
} from "./trackChangesModel";

export interface AnnotationInputs {
  baselineText: string | undefined;
  baselineReason: "entered" | "pinned" | "head" | undefined;
  spans: AuthorSpan[];
  proposals: ProposalView[];
  enabled: boolean;
}

/**
 * Inject the sentinel vocabulary into `text` block-by-block (INV-23: code/
 * mermaid fences stay atomic — never sentinel-injected, never word-refined) so
 * the delimiter-run + block-marker safety `injectSentinels` carries applies
 * PER BLOCK, not just at the very top of the whole document (its clamps are
 * anchored to a block's own start, matching every other caller — `colorByAuthor`
 * is always invoked per block too).
 */
function injectAllSentinels(text: string, marks: TaggedSpan[]): string {
  if (marks.length === 0) return text;
  const blocks = splitBlocksWithRanges(text);
  if (blocks.length === 0) return text;
  let out = "";
  let cursor = 0;
  for (const b of blocks) {
    out += text.slice(cursor, b.start); // preserve inter-block gaps (blank lines) verbatim
    if (b.type === "prose") {
      const blockMarks = marks.filter((m) => m.end > b.start && m.start < b.end);
      out += blockMarks.length ? injectSentinels(b.raw, b.start, blockMarks) : b.raw;
    } else {
      out += b.raw; // INV-23: fences are never annotated
    }
    cursor = b.end;
  }
  out += text.slice(cursor); // trailing content after the last block
  return out;
}

/**
 * Pure: markdown source -> sentinel-annotated source (authorship + change +
 * proposal marks vs the F6/F7 baseline). Skips everything when annotations are
 * disabled, or (the #48 rule) once the baseline was just PINNED with zero
 * LANDED diff — pending F4 proposals are unaffected by that gate (matching the
 * shipped preview's "pin→clean, proposals still show", #48/INV-33).
 */
export function annotateSource(src: string, inputs: AnnotationInputs): string {
  if (!inputs.enabled) return src;
  const proposals = inputs.proposals ?? [];
  const baselineText = inputs.baselineText ?? src; // no baseline → no change-marks
  const landedText = landedTextOf(src, proposals);
  const pinnedZeroDiff = inputs.baselineReason === "pinned" && landedText === baselineText;

  const insSpans: TaggedSpan[] = [];
  const deletions: { at: number; text: string }[] = [];
  const appliedRanges: { start: number; end: number }[] = [];

  // (c) pending F4 proposals already sit applied in `src` (F12 optimistic
  // apply, INV-48) — mark the applied range ins-<author> (`p.author` — a
  // proposal defaults to "claude" per `ProposalView.author`, but a human can
  // also be the proposer, e.g. a human-authored suggestion routed through the
  // same F4 pending-proposal path) and resurface the pre-apply original as a
  // deletion at the same point.
  for (const p of proposals) {
    if (p.anchorStart === null || p.anchorEnd === null) continue;
    appliedRanges.push({ start: p.anchorStart, end: p.anchorEnd });
    if (p.anchorEnd > p.anchorStart) insSpans.push({ start: p.anchorStart, end: p.anchorEnd, tag: p.author ?? "claude" });
    const original = p.original ?? p.replaced;
    if (original) deletions.push({ at: p.anchorStart, text: original });
  }

  // (a)/(b) committed changes-since-baseline. Each hunk (wordEditHunks(src,
  // baseline) — start/end in SRC coords, `replacement` = the baseline-side text
  // for that span) yields (a) insertion marks from the F3 `spans` CLIPPED to the
  // hunk's changed range — a hunk may straddle more than one author (e.g. two
  // adjacent words typed by different authors with nothing committed between
  // them), so spans are clipped individually, never the whole hunk tagged with
  // one author — and (b) the baseline's dropped text reinserted at the hunk
  // start. A hunk owned by an applied pending proposal is already marked above.
  if (!pinnedZeroDiff && landedText !== baselineText) {
    for (const h of wordEditHunks(src, baselineText)) {
      if (appliedRanges.some((r) => h.start < r.end && h.end > r.start)) continue;
      for (const s of inputs.spans) {
        const start = Math.max(s.start, h.start);
        const end = Math.min(s.end, h.end);
        if (end > start) insSpans.push({ start, end, tag: s.author });
      }
      if (h.replacement) deletions.push({ at: h.start, text: h.replacement });
    }
  }

  if (insSpans.length === 0 && deletions.length === 0) return src;

  // Splice deletion reinsertions into the source high→low (so earlier offsets
  // stay valid), shifting every mark recorded so far and appending a "del" mark
  // over the freshly-inserted text.
  const marks: TaggedSpan[] = insSpans.map((s) => ({ ...s }));
  let text = src;
  for (const d of [...deletions].sort((a, b) => b.at - a.at)) {
    text = text.slice(0, d.at) + d.text + text.slice(d.at);
    for (const m of marks) {
      if (m.start >= d.at) m.start += d.text.length;
      if (m.end >= d.at) m.end += d.text.length;
    }
    marks.push({ start: d.at, end: d.at + d.text.length, tag: "del" });
  }

  return injectAllSentinels(text, marks);
}

/** One `buildMermaidQueue` entry: a changed mermaid fence's own body (for
 * matching, `mermaidFenceBody` normalization — no trailing newline) paired
 * with its `diffMermaid`-augmented replacement source. `consumed` is reset by
 * the `render` wrapper (see `cowritingMarkdownItPlugin`) so the SAME entries
 * can back a re-render without a fresh parse. */
interface MermaidQueueEntry {
  readonly curBody: string;
  readonly source: string;
  consumed: boolean;
}

/**
 * Per-render queue of mermaid-fence overrides: one entry per `changed` atomic
 * mermaid op (Task 7 §2.6 parity: reuses `diffBlocks`' baseline/current block
 * pairing, the SAME pairing `renderReview`/`renderOp` used in the sunset
 * webview, so a changed mermaid fence is re-emitted through `diffMermaid`
 * exactly as it was there — INV-23 fences stay atomic, INV-29..31 styling
 * directives) that has a diffable before/current pair AND a successful
 * `diffMermaid` augmentation. `[]` when there is no baseline to diff against,
 * annotations are disabled, or the landed text equals the baseline (covers
 * the #48 pinned-zero-diff case too, since `diffBlocks` on identical text
 * yields no `changed` ops).
 *
 * Diffed against `landedText` (proposals reverted), matching the "committed
 * change since baseline" boundary `annotateSource`'s own `wordEditHunks(src,
 * baselineText)` gate already draws (a pending, not-yet-accepted proposal
 * isn't a landed change) — a fence a pending proposal's anchor overlaps is
 * left unaugmented, same as the rest of that gate.
 *
 * NOT positional: `src` (what markdown-it actually parses fences from) and
 * `landedText` (what this function diffs) can disagree on mermaid-fence COUNT
 * or ORDER — e.g. a pending proposal optimistically applied in `src` (F12,
 * INV-48) inserts a whole new mermaid fence that `landedText` (proposals
 * reverted) doesn't have, or `annotateSource`'s committed-deletion reinsert
 * splices an extra fence into the rendered source. `options.highlight`
 * (`cowritingMarkdownItPlugin`) matches each queue entry to a fence by BODY,
 * not by position — a fence with no matching entry renders verbatim (the
 * queue is a bag of `changed` fences to find, not a strict per-fence
 * override list), so misalignment degrades to verbatim exactly as intended,
 * never substitutes the wrong diagram's augmentation into the wrong fence.
 *
 * Pure, vscode-free (INV-6).
 */
function buildMermaidQueue(src: string, inputs: AnnotationInputs): MermaidQueueEntry[] {
  if (!inputs.enabled || inputs.baselineText === undefined) return [];
  const baselineText = inputs.baselineText;
  const landedText = landedTextOf(src, inputs.proposals ?? []);
  if (landedText === baselineText) return [];
  const entries: MermaidQueueEntry[] = [];
  for (const op of diffBlocks(baselineText, landedText)) {
    if (op.kind !== "changed" || op.block.type !== "mermaid") continue; // unchanged/added/removed: no override needed
    const curBody = mermaidFenceBody(op.block.raw);
    const beforeBody = op.before.type === "mermaid" ? mermaidFenceBody(op.before.raw) : null;
    if (curBody === null || beforeBody === null) continue;
    const result = diffMermaid(beforeBody, curBody);
    if (result.kind === "augmented") entries.push({ curBody, source: result.source, consumed: false }); // INV-30 fallback: verbatim
  }
  return entries;
}

/**
 * markdown-it plugin factory: the host supplies `AnnotationInputs` per render
 * (`host.inputsFor(env)`; `undefined` = no coediting context for this render —
 * source passed through untouched). Installs (a) a core rule BEFORE `normalize`
 * that swaps `state.src` for `annotateSource(...)`, and (b) a full-render
 * wrapper (see module docstring) that turns surviving sentinel pairs into
 * `cw-…` spans. Also teaches ```mermaid fences to render as `<pre class="mermaid">`
 * (Q4/Task 7 Step 2.6) via `options.highlight` — the same extension point the
 * built-in preview's own fence rule already calls for syntax highlighting — so
 * the bundled `media/preview-mermaid.js` previewScript has something to run
 * mermaid over (the built-in preview does not render mermaid on its own). A
 * CHANGED mermaid fence's source is re-emitted through `diffMermaid` first
 * (`buildMermaidQueue`, Task 7 §2.6 parity) so mermaid renders the intra-
 * diagram diff (INV-29..31), not just the current diagram. `options.highlight`
 * has no `env`/`state` parameter of its own, so the per-render queue is
 * threaded through a closure variable set by the core rule (which DOES see
 * `state.env`) and consumed BY BODY MATCH, not fence position (see
 * `buildMermaidQueue`), by `highlight`: each fence's `code` is matched against
 * unconsumed queue entries in order, the first match is consumed (so a later
 * fence with the same body still finds its own entry), and a fence with no
 * match renders verbatim WITHOUT consuming anything (so a later fence can
 * still match). markdown-it's `parse` (core rules) always completes before
 * `renderer.render` invokes `highlight`, and a single md instance never
 * renders concurrently, so the closure handoff itself is safe — but VS Code
 * may call `md.parse()` and `md.renderer.render()` separately and re-render
 * cached tokens without a fresh parse (the same caller pattern the module
 * docstring notes for scroll-sync), which would otherwise replay a queue whose
 * entries were already marked consumed by the prior render; the `render`
 * wrapper resets every entry's `consumed` flag so each render pass starts with
 * the full queue available again.
 */
export function cowritingMarkdownItPlugin(
  md: MarkdownIt,
  host: { inputsFor(env: unknown): AnnotationInputs | undefined },
): MarkdownIt {
  let mermaidQueue: MermaidQueueEntry[] = [];

  md.core.ruler.before("normalize", "cowriting-annotate", (state) => {
    const inputs = host.inputsFor(state.env);
    mermaidQueue = inputs ? buildMermaidQueue(state.src, inputs) : [];
    if (!inputs) return;
    state.src = annotateSource(state.src, inputs);
  });

  const baseRender = md.renderer.render.bind(md.renderer);
  md.renderer.render = (tokens, options, env) => {
    for (const entry of mermaidQueue) entry.consumed = false; // re-render-without-reparse safety net
    return sentinelsToSpans(baseRender(tokens, options, env), "ins");
  };

  const baseHighlight = md.options.highlight;
  md.options.highlight = (code, lang, attrs) => {
    if (lang?.trim().toLowerCase() === "mermaid") {
      const normalized = code.replace(/\n+$/, ""); // token.content keeps a trailing LF; mermaidFenceBody doesn't
      const entry = mermaidQueue.find((e) => !e.consumed && e.curBody === normalized);
      if (entry) entry.consumed = true;
      const legend = entry ? MERMAID_LEGEND : "";
      return `<pre class="mermaid" style="all: unset;">${escapeMermaidSource(entry?.source ?? code)}</pre>${legend}`;
    }
    return baseHighlight?.(code, lang, attrs) ?? md.utils.escapeHtml(code);
  };

  return md;
}

/** Minimal HTML-escape for a mermaid fence body (no markdown parsing — verbatim diagram source). */
function escapeMermaidSource(source: string): string {
  return source
    .replace(/&/g, "&amp;")
    .replace(/</g, "&lt;")
    .replace(/>/g, "&gt;")
    .replace(/\n+$/, "");
}