petal/web/src/components/Editor/SearchHighlight.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 { mapOffset } from './SuggestionHighlight'

// SearchHighlight powers the in-document Find & Replace bar. Like the suggestion
// layer it uses ProseMirror *decorations* (not stored marks), so matches are
// recomputed from the live document on every edit and never leave a stale
// highlight behind. Matches are found per-textblock (a query never spans a
// paragraph break), mirroring how the rest of the editor anchors text.

export interface Match {
  from: number
  to: number
}

interface PluginState {
  query: string
  caseSensitive: boolean
  matches: Match[]
  active: number // index into matches, or -1 when there are none
  decorations: DecorationSet
}

export const searchPluginKey = new PluginKey<PluginState>('petalSearch')

// findMatches collects every occurrence of `query` across the document's
// textblocks and maps each to an absolute ProseMirror range.
function findMatches(doc: PMNode, query: string, caseSensitive: boolean): Match[] {
  if (!query) return []
  const needle = caseSensitive ? query : query.toLowerCase()
  const matches: Match[] = []
  doc.descendants((node, pos) => {
    if (!node.isTextblock) return true
    const haystackRaw = node.textContent
    const haystack = caseSensitive ? haystackRaw : haystackRaw.toLowerCase()
    let idx = haystack.indexOf(needle)
    while (idx >= 0) {
      matches.push({
        from: mapOffset(node, pos, idx),
        to: mapOffset(node, pos, idx + query.length),
      })
      idx = haystack.indexOf(needle, idx + query.length)
    }
    return false // don't descend into inline children
  })
  return matches
}

function build(doc: PMNode, query: string, caseSensitive: boolean, preferred: number): PluginState {
  const matches = findMatches(doc, query, caseSensitive)
  const active = matches.length === 0 ? -1 : Math.max(0, Math.min(preferred, matches.length - 1))
  const decos = matches.map((m, i) =>
    Decoration.inline(m.from, m.to, {
      class: i === active ? 'petal-find-match petal-find-match-active' : 'petal-find-match',
    }),
  )
  return { query, caseSensitive, matches, active, decorations: DecorationSet.create(doc, decos) }
}

const EMPTY: PluginState = {
  query: '',
  caseSensitive: false,
  matches: [],
  active: -1,
  decorations: DecorationSet.empty,
}

// setSearch updates the query / case-sensitivity and recomputes matches. Passing
// an empty query clears the layer.
export function setSearch(
  state: EditorState,
  dispatch: (tr: Transaction) => void,
  query: string,
  caseSensitive: boolean,
) {
  dispatch(state.tr.setMeta(searchPluginKey, { kind: 'search', query, caseSensitive }))
}

// setActive moves the highlighted (current) match by index.
export function setActive(state: EditorState, dispatch: (tr: Transaction) => void, index: number) {
  dispatch(state.tr.setMeta(searchPluginKey, { kind: 'active', index }))
}

// clearSearch tears the layer down (on close).
export function clearSearch(state: EditorState, dispatch: (tr: Transaction) => void) {
  dispatch(state.tr.setMeta(searchPluginKey, { kind: 'clear' }))
}

// getSearchState reads the live plugin state (match list + active index) so the
// Find bar can render "n / m" and drive navigation/replacement.
export function getSearchState(state: EditorState): PluginState {
  return searchPluginKey.getState(state) ?? EMPTY
}

type Meta =
  | { kind: 'search'; query: string; caseSensitive: boolean }
  | { kind: 'active'; index: number }
  | { kind: 'clear' }

export const SearchHighlight = Extension.create({
  name: 'searchHighlight',

  addProseMirrorPlugins() {
    return [
      new Plugin<PluginState>({
        key: searchPluginKey,
        state: {
          init: () => EMPTY,
          apply(tr, value, _oldState, newState): PluginState {
            const meta = tr.getMeta(searchPluginKey) as Meta | undefined
            if (meta?.kind === 'search') {
              return build(newState.doc, meta.query, meta.caseSensitive, value.active < 0 ? 0 : value.active)
            }
            if (meta?.kind === 'active') {
              if (value.matches.length === 0) return value
              const active = ((meta.index % value.matches.length) + value.matches.length) % value.matches.length
              const decos = value.matches.map((m, i) =>
                Decoration.inline(m.from, m.to, {
                  class: i === active ? 'petal-find-match petal-find-match-active' : 'petal-find-match',
                }),
              )
              return { ...value, active, decorations: DecorationSet.create(newState.doc, decos) }
            }
            if (meta?.kind === 'clear') return EMPTY
            // Re-anchor on any document change so highlights track edits/replaces.
            if (tr.docChanged && value.query) {
              return build(newState.doc, value.query, value.caseSensitive, value.active)
            }
            return value
          },
        },
        props: {
          decorations(state) {
            return searchPluginKey.getState(state)?.decorations
          },
        },
      }),
    ]
  },
})