petal/web/src/components/Editor/SpellCheck.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'
import type { SpellChecker } from '../../hooks/useSpellChecker'

// SpellCheck renders browser-side nspell misspellings as ProseMirror
// decorations (a wavy red underline), recomputed from the live document on every
// change. Like the AI-suggestion layer it stores no marks — the underline is a
// pure overlay, so it never travels into saved content. English only: the word
// tokenizer matches Latin-letter runs, so CJK text (the user writes in both
// Mandarin and English) is simply never tokenized and never flagged.

export const spellPluginKey = new PluginKey<PluginState>('petalSpellCheck')

interface PluginState {
  checker: SpellChecker | null
  decorations: DecorationSet
}

// A word is a run of Latin letters with optional internal/edge apostrophes
// (don't, O'Brien). Anything else — digits, punctuation, CJK — terminates a run.
const WORD_RE = /[A-Za-z][A-Za-z']*/g

// isCheckable filters tokens we shouldn't flag: single letters and all-caps
// acronyms (NASA, USA), which dictionaries reliably miss and which read as noise
// when underlined.
function isCheckable(word: string): boolean {
  if (word.length < 2) return false
  if (word === word.toUpperCase()) return false
  return true
}

// strip leading/trailing apostrophes (e.g. a quoted 'word') so the dictionary
// lookup sees the bare token; returns the core plus how many chars were trimmed
// off the front (to re-anchor the decoration).
function coreOf(word: string): { core: string; lead: number } {
  const lead = word.match(/^'+/)?.[0].length ?? 0
  const trail = word.match(/'+$/)?.[0].length ?? 0
  return { core: word.slice(lead, word.length - trail), lead }
}

function eachMisspelling(
  doc: PMNode,
  checker: SpellChecker,
  visit: (from: number, to: number, word: string) => void,
) {
  doc.descendants((node, pos) => {
    if (!node.isTextblock) return true
    const text = node.textContent
    WORD_RE.lastIndex = 0
    let m: RegExpExecArray | null
    while ((m = WORD_RE.exec(text)) !== null) {
      const { core, lead } = coreOf(m[0])
      if (!isCheckable(core) || checker.correct(core)) continue
      const from = mapOffset(node, pos, m.index + lead)
      const to = mapOffset(node, pos, m.index + lead + core.length)
      visit(from, to, core)
    }
    return false // never descend into a textblock's inline children
  })
}

function buildDecorations(doc: PMNode, checker: SpellChecker, cursor: number): DecorationSet {
  const decos: Decoration[] = []
  eachMisspelling(doc, checker, (from, to, _word) => {
    // Don't flag the word the caret currently sits in — it's mid-typing, and a
    // red underline appearing under the cursor on every keystroke is jittery.
    if (cursor >= from && cursor <= to) return
    decos.push(Decoration.inline(from, to, { class: 'petal-misspelling', 'data-misspelling': '' }))
  })
  return DecorationSet.create(doc, decos)
}

// wordAt resolves the misspelling token under a ProseMirror position (from a
// click), returning its range + text so the card can offer corrections and the
// replacement can target the exact span — robust to duplicate words anywhere
// else in the document. Returns null if the position isn't inside a Latin word.
export function wordAt(doc: PMNode, pos: number): { from: number; to: number; word: string } | null {
  let found: { from: number; to: number; word: string } | null = null
  doc.descendants((node, nodePos) => {
    if (found) return false
    if (!node.isTextblock) return true
    if (pos <= nodePos || pos >= nodePos + node.nodeSize) return false
    const text = node.textContent
    WORD_RE.lastIndex = 0
    let m: RegExpExecArray | null
    while ((m = WORD_RE.exec(text)) !== null) {
      const { core, lead } = coreOf(m[0])
      if (!core) continue
      const from = mapOffset(node, nodePos, m.index + lead)
      const to = mapOffset(node, nodePos, m.index + lead + core.length)
      if (pos >= from && pos <= to) {
        found = { from, to, word: core }
        break
      }
    }
    return false
  })
  return found
}

// setSpellChecker pushes the (possibly null) checker into the plugin, rebuilding
// decorations immediately against the current document.
export function setSpellChecker(
  state: EditorState,
  dispatch: (tr: Transaction) => void,
  checker: SpellChecker | null,
) {
  dispatch(state.tr.setMeta(spellPluginKey, checker ?? null))
}

export const SpellCheck = Extension.create({
  name: 'spellCheck',

  addProseMirrorPlugins() {
    return [
      new Plugin<PluginState>({
        key: spellPluginKey,
        state: {
          init: () => ({ checker: null, decorations: DecorationSet.empty }),
          apply(tr, value, _oldState, newState) {
            const meta = tr.getMeta(spellPluginKey) as SpellChecker | null | undefined
            const checker = meta !== undefined ? meta : value.checker
            if (!checker) return { checker: null, decorations: DecorationSet.empty }
            // Rebuild on a checker swap, a doc edit, or a caret move (so the word
            // you just left gets re-evaluated and the new caret word is exempt).
            if (meta !== undefined || tr.docChanged || tr.selectionSet) {
              return { checker, decorations: buildDecorations(newState.doc, checker, newState.selection.head) }
            }
            return { checker, decorations: value.decorations }
          },
        },
        props: {
          decorations(state) {
            return spellPluginKey.getState(state)?.decorations
          },
        },
      }),
    ]
  },
})