llm-in-text/src/plugins/hiddenTextPlugin.ts

import { createApp, h, reactive } from 'vue'
import { Plugin, PluginKey } from '@milkdown/prose/state'
import { $node, $prose, $remark, $view } from '@milkdown/kit/utils'
import type { Node as ProseNode } from '@milkdown/prose/model'
import type { EditorView, NodeView } from '@milkdown/prose/view'
import HiddenTextCrepe from '../components/HiddenTextCrepe.vue'
import { HIDDEN_TEXT_NODE_TYPE, parseHiddenTextAt, serializeHiddenTextSyntax, splitTextWithHiddenSyntax } from '../utils/hiddenText.js'

const HIDDEN_TEXT_INPUT_PLUGIN_KEY = new PluginKey('milkdown-hidden-text-input')
const HIDDEN_TEXT_INPUT_META = 'hidden-text-input-meta'

function transformHiddenTextChildren(node: any) {
  if (!node || !Array.isArray(node.children)) return

  const nextChildren: any[] = []

  for (const child of node.children) {
    if (child?.type === 'text' && typeof child.value === 'string') {
      const segments = splitTextWithHiddenSyntax(child.value)
      if (segments) {
        nextChildren.push(...segments)
        continue
      }
    }

    transformHiddenTextChildren(child)
    nextChildren.push(child)
  }

  node.children = nextChildren
}

function findHiddenTextReplacements(doc: ProseNode) {
  const replacements: Array<{ from: number; to: number; displayed: string; hidden: string; marks: ProseNode['marks'] }> = []

  doc.descendants((node, pos, parent) => {
    if (!node.isText || !node.text) return true
    if (parent?.type.spec.code) return true
    if (node.marks.some((mark) => mark.type.spec.code)) return true

    let index = 0
    while (index < node.text.length) {
      const match = parseHiddenTextAt(node.text, index)
      if (!match) {
        index += 1
        continue
      }

      replacements.push({
        from: pos + match.start,
        to: pos + match.end,
        displayed: match.displayed,
        hidden: match.hidden,
        marks: node.marks,
      })

      index = match.end
    }

    return true
  })

  return replacements
}

class HiddenTextNodeView implements NodeView {
  node: ProseNode
  view: EditorView
  getPos: (() => number) | boolean
  dom: HTMLElement
  app: ReturnType<typeof createApp> | null = null
  props: Record<string, any>

  constructor(node: ProseNode, view: EditorView, getPos: (() => number) | boolean) {
    this.node = node
    this.view = view
    this.getPos = getPos
    this.dom = document.createElement('span')
    this.dom.className = 'hidden-text-node-view'
    this.dom.setAttribute('contenteditable', 'false')
    this.props = reactive({
      displayed: node.attrs.displayed,
      hidden: node.attrs.hidden,
      expanded: node.attrs.expanded,
      updateTexts: ({ displayed, hidden }: { displayed: string; hidden: string }) => {
        this.updateAttrs({ displayed, hidden })
      },
      updateExpanded: (expanded: boolean) => this.updateAttrs({ expanded }),
    })
    this.mount()
  }

  mount() {
    this.app = createApp({
      render: () => h(HiddenTextCrepe, this.props),
    })
    this.app.mount(this.dom)
  }

  getPosValue() {
    if (typeof this.getPos === 'function') {
      try {
        const pos = this.getPos()
        if (typeof pos === 'number') return pos
      } catch {
        // Fall back to DOM-based resolution below.
      }
    }

    try {
      const pos = this.view.posAtDOM(this.dom, 0)
      return typeof pos === 'number' ? pos : undefined
    } catch {
      return undefined
    }
  }

  updateAttrs(patch: Record<string, any>) {
    const pos = this.getPosValue()
    if (pos === undefined) return
    const nextAttrs = { ...this.node.attrs, ...patch }
    const nextNode = this.node.type.create(nextAttrs, undefined, this.node.marks)
    this.view.dispatch(this.view.state.tr.replaceWith(pos, pos + this.node.nodeSize, nextNode))
  }

  update(node: ProseNode) {
    if (node.type !== this.node.type) return false
    this.node = node
    this.props.displayed = node.attrs.displayed
    this.props.hidden = node.attrs.hidden
    this.props.expanded = node.attrs.expanded
    return true
  }

  stopEvent(event: Event) {
    const target = event.target as Node | null
    return Boolean(target && this.dom.contains(target))
  }

  ignoreMutation() {
    return true
  }

  destroy() {
    this.app?.unmount()
    this.app = null
  }
}

export const hiddenTextRemark = $remark('hiddenTextRemark', () => () => {
  return (tree: any) => {
    transformHiddenTextChildren(tree)
  }
})

export const hiddenTextNode = $node(HIDDEN_TEXT_NODE_TYPE, () => ({
  group: 'inline',
  inline: true,
  atom: true,
  selectable: true,
  draggable: false,
  attrs: {
    displayed: { default: '' },
    hidden: { default: '' },
    expanded: { default: false },
  },
  parseDOM: [
    {
      tag: 'span[data-hidden-text="true"]',
      getAttrs: (dom) => ({
        displayed: (dom as HTMLElement).getAttribute('data-hidden-display') || '',
        hidden: (dom as HTMLElement).getAttribute('data-hidden-value') || '',
        expanded: ((dom as HTMLElement).getAttribute('data-hidden-expanded') || '') === 'true',
      }),
    },
  ],
  toDOM: (node) => [
    'span',
    {
      'data-hidden-text': 'true',
      'data-hidden-display': node.attrs.displayed,
      'data-hidden-value': node.attrs.hidden,
      'data-hidden-expanded': String(Boolean(node.attrs.expanded)),
    },
  ],
  parseMarkdown: {
    match: (node) => node.type === HIDDEN_TEXT_NODE_TYPE,
    runner: (state, node, type) => {
      state.addNode(type, {
        displayed: String(node.displayed || ''),
        hidden: String(node.hidden || ''),
        expanded: false,
      })
    },
  },
  toMarkdown: {
    match: (node) => node.type.name === HIDDEN_TEXT_NODE_TYPE,
    runner: (state, node) => {
      state.addNode('text', undefined, serializeHiddenTextSyntax(node.attrs.displayed, node.attrs.hidden))
    },
  },
  leafText: (node) => serializeHiddenTextSyntax(node.attrs.displayed, node.attrs.hidden),
}))

export const hiddenTextView = $view(hiddenTextNode, () => {
  return (node, view, getPos) => new HiddenTextNodeView(node, view, getPos)
})

export const hiddenTextInputPlugin = $prose(() => {
  return new Plugin({
    key: HIDDEN_TEXT_INPUT_PLUGIN_KEY,
    appendTransaction: (transactions, _oldState, newState) => {
      if (!transactions.some((transaction) => transaction.docChanged)) return null
      if (transactions.some((transaction) => transaction.getMeta(HIDDEN_TEXT_INPUT_META))) return null

      const hiddenTextType = newState.schema.nodes[HIDDEN_TEXT_NODE_TYPE]
      if (!hiddenTextType) return null

      const replacements = findHiddenTextReplacements(newState.doc)
      if (replacements.length === 0) return null

      let tr = newState.tr

      for (let index = replacements.length - 1; index >= 0; index -= 1) {
        const replacement = replacements[index]
        tr = tr.replaceWith(
          replacement.from,
          replacement.to,
          hiddenTextType.create(
            {
              displayed: replacement.displayed,
              hidden: replacement.hidden,
              expanded: false,
            },
            undefined,
            replacement.marks
          )
        )
      }

      if (!tr.docChanged) return null
      tr.setMeta(HIDDEN_TEXT_INPUT_META, true)
      return tr
    },
  })
})