petal/web/src/components/Editor/SuggestionHighlight.ts

import { Extension } from '@tiptap/core'
import { Plugin, PluginKey } from '@tiptap/pm/state'
import type { EditorState, Transaction } from '@tiptap/pm/state'
import { Decoration, DecorationSet } from '@tiptap/pm/view'
import type { Node as PMNode } from '@tiptap/pm/model'
import type { Suggestion } from '../../api/client'

// SuggestionHighlight renders LLM suggestions as ProseMirror *decorations*, not
// stored marks. Decorations are ephemeral overlays recomputed from the live
// document on every change, which is exactly the string-anchoring the spec
// requires (Note #6): each suggestion is re-located by matching its `original`
// text in the current doc, so edits made while a checkpoint is in flight never
// leave a highlight stranded on stale coordinates.

export const suggestionPluginKey = new PluginKey<PluginState>('petalSuggestions')

interface PluginState {
  suggestions: Suggestion[]
  // The suggestion currently emphasized from its margin card / a hover, or null.
  // Carried in plugin state (not an imperative DOM class) so it survives the
  // decoration repaints that fire on every document change.
  activeId: string | null
  decorations: DecorationSet
}

// Meta carried on a transaction to update the plugin: either a fresh suggestion
// list or a change to which suggestion is emphasized.
type SuggestionMeta = { suggestions: Suggestion[] } | { activeId: string | null }

// mapOffset converts a character offset within a textblock's flattened text into
// an absolute ProseMirror position, accounting for inline atoms (e.g. hard
// breaks) that occupy a position but contribute no text.
export function mapOffset(block: PMNode, blockPos: number, targetOffset: number): number {
  let textOffset = 0
  let pmPos = blockPos + 1 // inline content starts just inside the block
  let result = pmPos
  let done = false
  block.forEach((child) => {
    if (done) return
    const len = child.isText ? (child.text?.length ?? 0) : 0
    if (textOffset + len >= targetOffset) {
      result = pmPos + (targetOffset - textOffset)
      done = true
    } else {
      textOffset += len
      pmPos += child.nodeSize
    }
  })
  if (!done) result = pmPos
  return result
}

// foldTypography canonicalizes the typographic variants that the editor's
// Typography input rules introduce (curly quotes, em/en dashes, single-glyph
// ellipsis) and the unicode spaces a paste can carry. The model echoes a
// suggestion's `original` from the same text but often normalizes these back to
// plain ASCII — straight quotes, `--`, `...` — so an exact match would miss and
// the suggestion could be neither highlighted nor applied.
//
// It returns the folded string plus `map`, where map[k] is the source index at
// which folded character k begins. A multi-character source run (`...`, `--`)
// folds to one glyph, so lengths differ; map projects a match in folded space
// back onto real document offsets. map has a trailing sentinel (= source length)
// so map[start + folded.length] yields the offset just past the match.
export function foldTypography(text: string): { folded: string; map: number[] } {
  let folded = ''
  const map: number[] = []
  let i = 0
  while (i < text.length) {
    // Multi-character ASCII sequences the Typography rules would have replaced.
    if (text.startsWith('...', i)) {
      folded += '…'
      map.push(i)
      i += 3
      continue
    }
    if (text[i] === '-' && text[i + 1] === '-') {
      folded += '—'
      map.push(i)
      i += 2
      continue
    }
    folded += foldChar(text[i])
    map.push(i)
    i++
  }
  map.push(text.length)
  return { folded, map }
}

// foldChar maps a single character onto its canonical form (length-preserving).
function foldChar(c: string): string {
  switch (c) {
    case '‘':
    case '’':
    case '‚':
    case '‛':
      return "'"
    case '“':
    case '”':
    case '„':
    case '‟':
      return '"'
    case '–': // en dash → em dash, the canonical form `--` also folds to
      return '—'
    case ' ': // non-breaking space
    case ' ': // thin space
    case ' ': // narrow no-break space
      return ' '
    default:
      return c
  }
}

// findRange locates the first occurrence of `search` within a single textblock
// and returns its ProseMirror range, or null if the string isn't present (the
// user may have edited or removed it since the checkpoint ran). Matching is done
// in canonical typographic space (see foldTypography) so curly-quote/dash/ellipsis
// differences between the model's echo and the live document don't strand the
// anchor. Exported so the accept flow can resolve the same span to replace.
export function findRange(doc: PMNode, search: string): { from: number; to: number } | null {
  const needle = foldTypography(search).folded
  if (!needle) return null
  let result: { from: number; to: number } | null = null
  doc.descendants((node, pos) => {
    if (result) return false
    if (!node.isTextblock) return true // keep descending to the textblock
    const { folded, map } = foldTypography(node.textContent)
    const idx = folded.indexOf(needle)
    if (idx >= 0) {
      result = {
        from: mapOffset(node, pos, map[idx]),
        to: mapOffset(node, pos, map[idx + needle.length]),
      }
    }
    return false // never descend into a textblock's inline children
  })
  return result
}

function buildDecorations(doc: PMNode, suggestions: Suggestion[], activeId: string | null): DecorationSet {
  const decos: Decoration[] = []
  for (const s of suggestions) {
    const range = findRange(doc, s.original)
    if (!range) continue
    const active = s.id === activeId ? ' petal-suggestion-active' : ''
    decos.push(
      Decoration.inline(range.from, range.to, {
        class: `petal-suggestion petal-suggestion-${s.type}${active}`,
        'data-suggestion-id': s.id,
      }),
    )
  }
  return DecorationSet.create(doc, decos)
}

// setSuggestions pushes a new suggestion list into the plugin. Decorations are
// rebuilt against the current document immediately.
export function setSuggestions(state: EditorState, dispatch: (tr: Transaction) => void, suggestions: Suggestion[]) {
  const tr = state.tr.setMeta(suggestionPluginKey, { suggestions } satisfies SuggestionMeta)
  dispatch(tr)
}

// setActiveSuggestion marks one suggestion (or none) as emphasized, rebuilding
// the decorations so the flagged text gets the active wash. Driven by the margin
// rail's hover/selection so the card↔text link reads both ways.
export function setActiveSuggestion(
  state: EditorState,
  dispatch: (tr: Transaction) => void,
  activeId: string | null,
) {
  if (suggestionPluginKey.getState(state)?.activeId === activeId) return
  dispatch(state.tr.setMeta(suggestionPluginKey, { activeId } satisfies SuggestionMeta))
}

export const SuggestionHighlight = Extension.create({
  name: 'suggestionHighlight',

  addProseMirrorPlugins() {
    return [
      new Plugin<PluginState>({
        key: suggestionPluginKey,
        state: {
          init: () => ({ suggestions: [], activeId: null, decorations: DecorationSet.empty }),
          apply(tr, value, _oldState, newState) {
            const meta = tr.getMeta(suggestionPluginKey) as SuggestionMeta | undefined
            if (meta && 'suggestions' in meta) {
              return {
                suggestions: meta.suggestions,
                activeId: value.activeId,
                decorations: buildDecorations(newState.doc, meta.suggestions, value.activeId),
              }
            }
            if (meta && 'activeId' in meta) {
              return {
                suggestions: value.suggestions,
                activeId: meta.activeId,
                decorations: buildDecorations(newState.doc, value.suggestions, meta.activeId),
              }
            }
            // On any document change, re-anchor by string against the new doc.
            if (tr.docChanged) {
              return {
                suggestions: value.suggestions,
                activeId: value.activeId,
                decorations: buildDecorations(newState.doc, value.suggestions, value.activeId),
              }
            }
            return value
          },
        },
        props: {
          decorations(state) {
            return suggestionPluginKey.getState(state)?.decorations
          },
        },
      }),
    ]
  },
})