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

import { createApp, reactive } from 'vue'
import { serializerCtx } from '@milkdown/kit/core'
import { $node, $remark, $view } from '@milkdown/kit/utils'
import type { Node as ProseNode, Schema } from '@milkdown/prose/model'
import type { EditorView, NodeView } from '@milkdown/prose/view'
import DocBlockCrepe from '../components/DocBlockCrepe.vue'
import {
  DOC_BLOCK_FENCE_LANG,
  DOC_BLOCK_NODE_TYPE,
  DOC_CONTEXT_LIMIT,
  buildLegacyDocBlock,
  buildDocContextFence,
  normalizeDocType,
  parseLegacyDocBlock,
  parseDocBlockValue,
  stripDocBlockMarkdown,
} from '../utils/docBlock.js'

function serializeRangeToMarkdown(
  doc: ProseNode,
  from: number,
  to: number,
  schema: Schema,
  serializer: (content: ProseNode) => string
): string {
  if (from >= to) return ''
  const slice = doc.slice(from, to)
  if (slice.content.size <= 0) return ''
  const sliceDoc = schema.topNodeType.createAndFill(undefined, slice.content)
  return sliceDoc ? serializer(sliceDoc) : doc.textBetween(from, to, '\n', '\n')
}

function buildDocContext(doc: ProseNode, excludePos?: number) {
  const blocks: string[] = []
  doc.descendants((node, pos) => {
    if (node.type.name !== DOC_BLOCK_NODE_TYPE) return true
    if (excludePos !== undefined && pos === excludePos) return false
    blocks.push(
      buildDocContextFence({
        docType: node.attrs.docType,
        content: node.attrs.content,
      })
    )
    return false
  })
  return blocks.join('\n\n')
}

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

  constructor(node: ProseNode, view: EditorView, getPos: () => number | undefined, serializer: (content: ProseNode) => string) {
    this.node = node
    this.view = view
    this.getPos = getPos
    this.serializer = serializer
    this.dom = document.createElement('div')
    this.dom.className = 'doc-block-node-view'
    this.props = reactive({
      docType: node.attrs.docType,
      docName: node.attrs.docName,
      uploadTime: node.attrs.uploadTime,
      content: node.attrs.content,
      collapsed: node.attrs.collapsed,
      onUpdateContent: (content: string) => this.updateAttrs({ content }),
      onUpdateCollapsed: (collapsed: boolean) => this.updateAttrs({ collapsed }),
      onDelete: () => this.deleteNode(),
      resolveSuggestionRequest: (payload: { prefix: string; suffix: string; languageId: string }) => this.resolveSuggestionRequest(payload),
    })
    this.mount()
  }

  mount() {
    this.app = createApp(DocBlockCrepe, this.props)
    this.app.mount(this.dom)
  }

  getPosValue() {
    const pos = this.getPos()
    return typeof pos === 'number' ? pos : 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))
  }

  deleteNode() {
    const pos = this.getPosValue()
    if (pos === undefined) return
    const tr = this.view.state.tr.delete(pos, pos + this.node.nodeSize).scrollIntoView()
    this.view.dispatch(tr)
    this.view.focus()
  }

  resolveSuggestionRequest(payload: { prefix: string; suffix: string; languageId: string }) {
    const pos = this.getPosValue()
    if (pos === undefined) return payload
    const doc = this.view.state.doc
    const schema = this.view.state.schema
    const before = stripDocBlockMarkdown(serializeRangeToMarkdown(doc, 0, pos, schema, this.serializer))
    const after = stripDocBlockMarkdown(serializeRangeToMarkdown(doc, pos + this.node.nodeSize, doc.content.size, schema, this.serializer))
    const docContext = buildDocContext(doc, pos)
    const mergedPrefix = [docContext, before, payload.prefix].filter(Boolean).join('\n\n')
    const mergedSuffix = [payload.suffix, after].filter(Boolean).join('\n\n')
    if (mergedPrefix.length + mergedSuffix.length > DOC_CONTEXT_LIMIT) {
      return {
        prefix: mergedPrefix.slice(0, DOC_CONTEXT_LIMIT),
        suffix: '',
        languageId: payload.languageId,
        blocked: true,
      }
    }
    return {
      prefix: mergedPrefix,
      suffix: mergedSuffix,
      languageId: payload.languageId,
      blocked: false,
    }
  }

  update(node: ProseNode) {
    if (node.type !== this.node.type) return false
    this.node = node
    this.props.docType = node.attrs.docType
    this.props.docName = node.attrs.docName
    this.props.uploadTime = node.attrs.uploadTime
    this.props.content = node.attrs.content
    this.props.collapsed = node.attrs.collapsed
    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
  }
}

function visitChildren(node: any, visitor: (child: any) => any) {
  if (!node || !Array.isArray(node.children)) return
  node.children = node.children.map((child: any) => {
    const next = visitor(child)
    if (next && next !== child) return next
    visitChildren(child, visitor)
    return child
  })
}

export const docBlockRemark = $remark('docBlockRemark', () => () => {
  return (tree: any) => {
    visitChildren(tree, (node) => {
      if (node?.type === 'code' && node.lang === DOC_BLOCK_FENCE_LANG) {
        return {
          type: 'docBlock',
          value: String(node.value || ''),
          sourceType: 'code',
        }
      }
      if (node?.type === 'html' && typeof node.value === 'string' && node.value.includes('<doc_type=')) {
        return {
          type: 'docBlock',
          value: String(node.value || ''),
          sourceType: 'html',
        }
      }
      return node
    })
  }
})

export const docBlockNode = $node(DOC_BLOCK_NODE_TYPE, () => ({
  group: 'block',
  atom: true,
  isolating: true,
  selectable: true,
  draggable: false,
  marks: '',
  attrs: {
    docType: { default: 'txt' },
    docName: { default: 'document.txt' },
    uploadTime: { default: '' },
    content: { default: '' },
    collapsed: { default: false },
  },
  parseDOM: [
    {
      tag: 'div[data-doc-block="true"]',
      getAttrs: (dom) => ({
        docType: normalizeDocType((dom as HTMLElement).getAttribute('data-doc-type') || ''),
        docName: (dom as HTMLElement).getAttribute('data-doc-name') || 'document.txt',
        uploadTime: (dom as HTMLElement).getAttribute('data-doc-upload-time') || '',
        collapsed: ((dom as HTMLElement).getAttribute('data-doc-collapsed') || '') === 'true',
        content: '',
      }),
    },
  ],
  toDOM: (node) => [
    'div',
    {
      'data-doc-block': 'true',
      'data-doc-type': node.attrs.docType,
      'data-doc-name': node.attrs.docName,
      'data-doc-upload-time': node.attrs.uploadTime,
      'data-doc-collapsed': String(Boolean(node.attrs.collapsed)),
    },
  ],
  parseMarkdown: {
    match: (node) => node.type === 'docBlock',
    runner: (state, node, type) => {
      const attrs = node.sourceType === 'code'
        ? parseDocBlockValue(String(node.value || ''))
        : parseLegacyDocBlock(String(node.value || ''))
      if (!attrs) return
      state.addNode(type, attrs)
    },
  },
  toMarkdown: {
    match: (node) => node.type.name === DOC_BLOCK_NODE_TYPE,
    runner: (state, node) => {
      state.addNode('html', undefined, buildLegacyDocBlock(node.attrs))
    },
  },
}))

export const docBlockView = $view(docBlockNode, (ctx) => {
  const serializer = ctx.get(serializerCtx)
  return (node, view, getPos) => new DocBlockNodeView(node, view, getPos, serializer)
})

export function buildDocContextFromDoc(doc: ProseNode, excludePos?: number) {
  return buildDocContext(doc, excludePos)
}