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

import { createApp, h, reactive } from 'vue'
import { parserCtx, serializerCtx } from '@milkdown/kit/core'
import { $ctx, $node, $prose, $remark, $view } from '@milkdown/kit/utils'
import { Plugin, PluginKey, Selection } from '@milkdown/prose/state'
import { type Node as ProseNode, Slice, type Schema } from '@milkdown/prose/model'
import { Decoration, DecorationSet, type EditorView, type NodeView } from '@milkdown/prose/view'
import ProBlockCrepe from '../components/ProBlockCrepe.vue'
import { extractDocBlockContextFromMarkdown } from '../utils/docBlock.js'
import { extractTextFromOCR, getOcrCache } from '../utils/ocrCache'
import { PRO_BLOCK_NODE_TYPE, PRO_DISPLAY_LABEL, PRO_TRIGGER_TEXT, parseProBlockSyntax, serializeProBlockSyntax } from '../utils/proBlock.js'

const PRO_BLOCK_INPUT_PLUGIN_KEY = new PluginKey('milkdown-pro-block-input')
const PRO_BLOCK_HIGHLIGHT_PLUGIN_KEY = new PluginKey<DecorationSet>('milkdown-pro-block-highlight')
const PRO_BLOCK_INPUT_META = 'pro-block-input-meta'
const PRO_CONTEXT_LIMIT = 32 * 1024
const FALLBACK_BLOCK_SEPARATOR = '\n\n'
const FALLBACK_LEAF_TEXT = '\n'
const IMAGE_NODE_TYPES = new Set(['image', 'image-block', 'imageBlock'])

interface ProBlockConfig {
  fetchSuggestionStream: (payload: {
    prefix: string
    suffix: string
    languageId: string
    signal?: AbortSignal
    model?: string
    temperature?: number
    onChunk?: (chunk: string) => void
  }) => Promise<string>
  getProModel: () => string
  showError: (message: string) => void
  isThinkingActive: () => boolean
  setThinkingActive: (value: boolean) => void
}

interface ProRequestPayload {
  prefix: string
  suffix: string
  languageId: string
  blocked: boolean
}

function serializeRangeToMarkdown(
  doc: ProseNode,
  from: number,
  to: number,
  schema: Schema,
  serializer: (content: ProseNode) => string
): string {
  if (from >= to) return ''
  const fallback = doc.textBetween(from, to, FALLBACK_BLOCK_SEPARATOR, FALLBACK_LEAF_TEXT)
  if (typeof serializer !== 'function') return fallback

  const slice = doc.slice(from, to)
  if (slice.content.size <= 0) return ''

  try {
    const sliceDoc = schema.topNodeType.createAndFill(undefined, slice.content)
    return sliceDoc ? serializer(sliceDoc) : fallback
  } catch {
    return fallback
  }
}

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

  node.children = node.children.map((child: any) => {
    const replacement = parseProParagraphNode(child)
    if (replacement) return replacement
    transformProBlockChildren(child)
    return child
  })
}

function parseProParagraphNode(node: any) {
  if (!node || node.type !== 'paragraph' || !Array.isArray(node.children)) return null
  if (node.children.length !== 1) return null

  const textNode = node.children[0]
  if (!textNode || textNode.type !== 'text' || typeof textNode.value !== 'string') return null

  const parsed = parseProBlockSyntax(textNode.value)
  if (!parsed) return null

  return {
    type: PRO_BLOCK_NODE_TYPE,
    content: parsed.content,
    autoStart: false,
  }
}

function findProBlockReplacements(doc: ProseNode) {
  const replacements: Array<{ from: number; to: number; content: string; autoStart: boolean }> = []

  doc.descendants((node, pos) => {
    if (node.type.name !== 'paragraph' || node.childCount !== 1) return true

    const firstChild = node.firstChild
    if (!firstChild?.isText) return true

    const parsed = parseProBlockSyntax(node.textContent)
    if (!parsed) return true

    replacements.push({
      from: pos,
      to: pos + node.nodeSize,
      content: parsed.content,
      autoStart: parsed.autoStart,
    })

    return false
  })

  return replacements
}

function buildOcrContext(doc: ProseNode) {
  const lines: string[] = []

  doc.descendants((node) => {
    if (!IMAGE_NODE_TYPES.has(node.type.name)) return true
    const src = typeof node.attrs?.src === 'string' ? node.attrs.src : ''
    if (!src) return true
    const ocrText = getOcrCache(src)
    if (!ocrText) return true

    const preview = extractTextFromOCR(ocrText, 120)
    if (!preview) return true

    const label = typeof node.attrs?.alt === 'string' && node.attrs.alt.trim() ? node.attrs.alt.trim() : 'image'
    lines.push(`![${label}](${src}) <OCR:${preview}>`)
    return true
  })

  return lines.join('\n')
}

function createHighlightDecorations(doc: ProseNode, from: number, to: number) {
  if (to <= from) return DecorationSet.empty

  const decorations = [
    Decoration.inline(from, to, { class: 'pro-block-accepted-highlight' }),
  ]

  doc.nodesBetween(from, to, (node, pos) => {
    if (!node.isBlock) return true
    const nodeFrom = pos
    const nodeTo = pos + node.nodeSize
    if (nodeFrom >= from && nodeTo <= to) {
      decorations.push(Decoration.node(nodeFrom, nodeTo, { class: 'pro-block-accepted-block' }))
    }
    return true
  })

  return DecorationSet.create(doc, decorations)
}

function replaceWithParsedMarkdownSlice(tr: any, from: number, to: number, parsedDoc: ProseNode) {
  if (!parsedDoc || parsedDoc.content.size <= 0) return null

  const parsedSlice = Slice.maxOpen(parsedDoc.content)
  if (!parsedSlice || parsedSlice.size <= 0) return null

  tr.replaceRange(from, to, parsedSlice)
  const startPos = tr.mapping.map(from, -1)
  const endPos = tr.mapping.map(to, 1)
  if (endPos <= startPos) return null
  return { from: startPos, to: endPos }
}

function insertProBlockNode(view: EditorView, autoStart: boolean) {
  const proBlockType = view.state.schema.nodes[PRO_BLOCK_NODE_TYPE]
  if (!proBlockType) return false

  const { from, to } = view.state.selection
  const blockNode = proBlockType.create({ content: '', autoStart })
  const tr = view.state.tr.replaceRangeWith(from, to, blockNode)
  const nextPos = Math.min(from + blockNode.nodeSize, tr.doc.content.size)
  tr.setSelection(Selection.near(tr.doc.resolve(nextPos), 1))
  view.dispatch(tr.scrollIntoView())
  view.focus()
  return true
}

function replaceParagraphWithProBlock(
  view: EditorView,
  from: number,
  to: number,
  content: string,
  autoStart: boolean
) {
  const proBlockType = view.state.schema.nodes[PRO_BLOCK_NODE_TYPE]
  if (!proBlockType) return false

  const blockNode = proBlockType.create({
    content,
    autoStart: autoStart && content === '',
  })
  const tr = view.state.tr.replaceWith(from, to, blockNode)
  const nextPos = Math.min(from + blockNode.nodeSize, tr.doc.content.size)
  tr.setMeta(PRO_BLOCK_INPUT_META, true)
  tr.setSelection(Selection.near(tr.doc.resolve(nextPos), 1))
  view.dispatch(tr.scrollIntoView())
  view.focus()
  return true
}

function tryHandleProTriggerTextInput(view: EditorView, insertedText: string) {
  if (!insertedText || !insertedText.includes(']')) return false

  const { state } = view
  const { $from, from, to } = state.selection
  const paragraph = $from.parent
  if (!paragraph || paragraph.type.name !== 'paragraph') return false
  if (!$from.sameParent(state.selection.$to)) return false

  const paragraphDepth = $from.depth
  const paragraphStart = $from.start(paragraphDepth)
  const startOffset = from - paragraphStart
  const endOffset = to - paragraphStart
  const nextText = `${paragraph.textContent.slice(0, startOffset)}${insertedText}${paragraph.textContent.slice(endOffset)}`
  const parsed = parseProBlockSyntax(nextText)
  if (!parsed?.autoStart || parsed.content) return false

  const blockFrom = $from.before(paragraphDepth)
  const blockTo = blockFrom + paragraph.nodeSize
  return replaceParagraphWithProBlock(view, blockFrom, blockTo, parsed.content, parsed.autoStart)
}

function isAbortError(error: unknown) {
  return Boolean(error && typeof error === 'object' && 'name' in error && (error as { name?: string }).name === 'AbortError')
}

export const proBlockConfigCtx = $ctx<ProBlockConfig, 'proBlockConfig'>({
  fetchSuggestionStream: async () => '',
  getProModel: () => '',
  showError: () => {},
  isThinkingActive: () => false,
  setThinkingActive: () => {},
}, 'proBlockConfig')

class ProBlockNodeView implements NodeView {
  node: ProseNode
  view: EditorView
  getPos: (() => number) | boolean
  dom: HTMLElement
  app: ReturnType<typeof createApp> | null = null
  props: Record<string, any>
  parser: (markdown: string) => Promise<ProseNode>
  serializer: (content: ProseNode) => string
  config: ProBlockConfig
  abortController: AbortController | null = null
  requestSeq = 0
  redoCount = 0
  destroyed = false
  highlightTimer: ReturnType<typeof setTimeout> | null = null

  constructor(
    node: ProseNode,
    view: EditorView,
    getPos: (() => number) | boolean,
    parser: (markdown: string) => Promise<ProseNode>,
    serializer: (content: ProseNode) => string,
    config: ProBlockConfig
  ) {
    this.node = node
    this.view = view
    this.getPos = getPos
    this.parser = parser
    this.serializer = serializer
    this.config = config
    this.dom = document.createElement('div')
    this.dom.className = 'pro-block-node-view'
    this.props = reactive({
      stage: node.attrs.content ? 'done' : 'idle',
      previewContent: '',
      savedContent: node.attrs.content || '',
      isBusy: false,
      title: PRO_DISPLAY_LABEL,
      activateAction: () => {
        void this.startThinking(false)
      },
      discardAction: () => {
        this.discardResult()
      },
      redoAction: () => {
        void this.startThinking(true)
      },
      acceptAction: () => {
        void this.acceptResult()
      },
    })
    this.mount()
    queueMicrotask(() => {
      if (!this.destroyed) {
        this.consumeAutoStart()
      }
    })
  }

  mount() {
    this.app = createApp({
      render: () => h(ProBlockCrepe, { ...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 {
        // Ignore and fall back to DOM lookup.
      }
    }

    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 }
    this.view.dispatch(this.view.state.tr.setNodeMarkup(pos, undefined, nextAttrs))
  }

  setStage(stage: 'idle' | 'thinking' | 'streaming' | 'done', previewContent = '') {
    this.props.stage = stage
    this.props.previewContent = previewContent
    this.props.isBusy = stage === 'thinking' || stage === 'streaming'
  }

  buildRequestPayload(): ProRequestPayload {
    const pos = this.getPosValue()
    if (pos === undefined) {
      return {
        prefix: '',
        suffix: '',
        languageId: 'markdown',
        blocked: true,
      }
    }

    const doc = this.view.state.doc
    const schema = this.view.state.schema
    const prefixMarkdown = serializeRangeToMarkdown(doc, 0, pos, schema, this.serializer)
      || doc.textBetween(0, pos, FALLBACK_BLOCK_SEPARATOR, FALLBACK_LEAF_TEXT)
    const suffixStart = Math.min(pos + this.node.nodeSize, doc.content.size)
    const suffixMarkdown = serializeRangeToMarkdown(doc, suffixStart, doc.content.size, schema, this.serializer)
      || doc.textBetween(suffixStart, doc.content.size, FALLBACK_BLOCK_SEPARATOR, FALLBACK_LEAF_TEXT)

    const ocrContext = buildOcrContext(doc)
    const docContext = extractDocBlockContextFromMarkdown(`${prefixMarkdown}\n\n${suffixMarkdown}`, 1600)
    const fullPrefix = [ocrContext, docContext, prefixMarkdown].filter(Boolean).join('\n\n')

    return {
      prefix: fullPrefix,
      suffix: suffixMarkdown,
      languageId: 'markdown',
      blocked: fullPrefix.length + suffixMarkdown.length > PRO_CONTEXT_LIMIT,
    }
  }

  consumeAutoStart() {
    if (!this.node.attrs.autoStart) return
    this.updateAttrs({ autoStart: false })
    void this.startThinking(false)
  }

  stopThinkingLock() {
    this.config.setThinkingActive(false)
  }

  abortActive(reason = 'abort') {
    if (!this.abortController) return
    this.abortController.abort(reason)
    this.abortController = null
    this.stopThinkingLock()
  }

  discardResult() {
    this.abortActive('discard')
    this.redoCount = 0
    this.props.savedContent = ''
    this.setStage('idle', '')
    this.updateAttrs({ content: '', autoStart: false })
  }

  async startThinking(isRedo: boolean) {
    if (this.props.isBusy) return
    const previousContent = this.node.attrs.content || ''
    let requestSeq = this.requestSeq
    let requestStarted = false

    try {
      if (this.config.isThinkingActive()) {
        this.config.showError('当前已有一个 PRO 思考正在进行，请先等待它完成。')
        return
      }

      const payload = this.buildRequestPayload()
      if (payload.blocked) {
        this.config.showError('PRO 模式上下文超过 32KB，无法继续思考。')
        return
      }

      requestSeq = this.requestSeq + 1
      this.requestSeq = requestSeq
      this.abortController = new AbortController()
      this.config.setThinkingActive(true)
      requestStarted = true
      this.setStage('thinking', '')

      const temperature = Math.min(1.2, 0.7 + (isRedo ? (this.redoCount + 1) * 0.2 : 0))

      const result = await this.config.fetchSuggestionStream({
        ...payload,
        signal: this.abortController.signal,
        model: this.config.getProModel(),
        temperature,
        onChunk: (chunk) => {
          if (this.destroyed || this.requestSeq !== requestSeq) return
          const nextContent = `${this.props.previewContent || ''}${chunk}`
          this.setStage('streaming', nextContent)
        },
      })

      if (this.destroyed || this.requestSeq !== requestSeq) return

      const content = String(result || this.props.previewContent || '').trim()
      this.redoCount = isRedo ? this.redoCount + 1 : 0
      this.props.savedContent = content
      this.setStage(content ? 'done' : 'idle', '')
      this.updateAttrs({ content, autoStart: false })
    } catch (error) {
      if (!isAbortError(error)) {
        this.config.showError(error instanceof Error ? error.message : 'PRO 模式思考失败')
      }

      this.props.savedContent = previousContent
      this.setStage(previousContent ? 'done' : 'idle', '')
    } finally {
      if (requestStarted && this.requestSeq === requestSeq) {
        this.abortController = null
        this.stopThinkingLock()
      }
    }
  }

  async acceptResult() {
    if (this.props.isBusy) return

    const source = String(this.props.savedContent || this.node.attrs.content || '').trim()
    if (!source) return

    const pos = this.getPosValue()
    if (pos === undefined) return

    const from = pos
    const to = pos + this.node.nodeSize
    const tr = this.view.state.tr
    let insertedRange: { from: number; to: number } | null = null

    try {
      const parsedDoc = await this.parser(source)
      insertedRange = replaceWithParsedMarkdownSlice(tr, from, to, parsedDoc)
    } catch {
      insertedRange = null
    }

    if (!insertedRange) {
      tr.insertText(source, from, to)
      insertedRange = { from, to: from + source.length }
    }

    const endPos = Math.min(insertedRange.to, tr.doc.content.size)
    tr.setSelection(Selection.near(tr.doc.resolve(endPos), 1))
    tr.setMeta(PRO_BLOCK_HIGHLIGHT_PLUGIN_KEY, { set: insertedRange })
    this.view.dispatch(tr.scrollIntoView())
    this.view.focus()

    if (this.highlightTimer) {
      clearTimeout(this.highlightTimer)
    }
    this.highlightTimer = setTimeout(() => {
      if (this.destroyed) return
      this.view.dispatch(this.view.state.tr.setMeta(PRO_BLOCK_HIGHLIGHT_PLUGIN_KEY, { clear: true }))
    }, 1000)
  }

  update(node: ProseNode) {
    if (node.type !== this.node.type) return false
    this.node = node

    if (!this.props.isBusy) {
      this.props.savedContent = node.attrs.content || ''
      this.props.previewContent = ''
      this.props.stage = node.attrs.content ? 'done' : 'idle'
    }

    return true
  }

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

  ignoreMutation() {
    return true
  }

  destroy() {
    this.destroyed = true
    this.abortActive('destroy')
    if (this.highlightTimer) {
      clearTimeout(this.highlightTimer)
      this.highlightTimer = null
    }
    this.app?.unmount()
    this.app = null
  }
}

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

export const proBlockNode = $node(PRO_BLOCK_NODE_TYPE, () => ({
  group: 'block',
  atom: true,
  isolating: true,
  selectable: true,
  draggable: false,
  marks: '',
  attrs: {
    content: { default: '' },
    autoStart: { default: false },
  },
  parseDOM: [
    {
      tag: 'div[data-pro-block="true"]',
      getAttrs: (dom) => ({
        content: (dom as HTMLElement).getAttribute('data-pro-content') || '',
        autoStart: ((dom as HTMLElement).getAttribute('data-pro-autostart') || '') === 'true',
      }),
    },
  ],
  toDOM: (node) => [
    'div',
    {
      'data-pro-block': 'true',
      'data-pro-content': node.attrs.content || '',
      'data-pro-autostart': String(Boolean(node.attrs.autoStart)),
    },
  ],
  parseMarkdown: {
    match: (node) => node.type === PRO_BLOCK_NODE_TYPE,
    runner: (state, node, type) => {
      state.addNode(type, {
        content: String(node.content || ''),
        autoStart: false,
      })
    },
  },
  toMarkdown: {
    match: (node) => node.type.name === PRO_BLOCK_NODE_TYPE,
    runner: (state, node) => {
      state.addNode('paragraph', [{
        type: 'text',
        value: serializeProBlockSyntax(node.attrs.content),
      }])
    },
  },
  leafText: (node) => serializeProBlockSyntax(node.attrs.content),
}))

export const proBlockView = $view(proBlockNode, (ctx) => {
  const parser = ctx.get(parserCtx)
  const serializer = ctx.get(serializerCtx)
  const config = ctx.get(proBlockConfigCtx.key)
  return (node, view, getPos) => new ProBlockNodeView(node, view, getPos, parser, serializer, config)
})

export const proBlockInputPlugin = $prose(() => {
  return new Plugin({
    key: PRO_BLOCK_INPUT_PLUGIN_KEY,
    props: {
      handleTextInput: (view, _from, _to, text) => {
        return tryHandleProTriggerTextInput(view, text)
      },
      handleKeyDown: (view, event) => {
        const pressed = (event.ctrlKey || event.metaKey) && event.shiftKey && event.key.toLowerCase() === 'p'
        if (!pressed) return false
        event.preventDefault()
        return insertProBlockNode(view, false)
      },
    },
    appendTransaction: (transactions, _oldState, newState) => {
      if (!transactions.some((transaction) => transaction.docChanged)) return null
      if (transactions.some((transaction) => transaction.getMeta(PRO_BLOCK_INPUT_META))) return null

      const proBlockType = newState.schema.nodes[PRO_BLOCK_NODE_TYPE]
      if (!proBlockType) return null

      const replacements = findProBlockReplacements(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,
          proBlockType.create({
            content: replacement.content,
            autoStart: replacement.autoStart && replacement.content === '',
          })
        )
      }

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

export const proBlockHighlightPlugin = $prose(() => {
  return new Plugin<DecorationSet>({
    key: PRO_BLOCK_HIGHLIGHT_PLUGIN_KEY,
    state: {
      init: () => DecorationSet.empty,
      apply: (tr, value) => {
        const meta = tr.getMeta(PRO_BLOCK_HIGHLIGHT_PLUGIN_KEY)
        if (meta?.clear) {
          return DecorationSet.empty
        }
        if (meta?.set) {
          return createHighlightDecorations(tr.doc, meta.set.from, meta.set.to)
        }
        return value.map(tr.mapping, tr.doc)
      },
    },
    props: {
      decorations: (state) => PRO_BLOCK_HIGHLIGHT_PLUGIN_KEY.getState(state),
    },
  })
})

export function insertProBlockAtSelection(view: EditorView, autoStart = false) {
  return insertProBlockNode(view, autoStart)
}

export { PRO_TRIGGER_TEXT }