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'

const FALLBACK_BLOCK_SEPARATOR = '\n\n'
const FALLBACK_LEAF_TEXT = '\n'
const CONTEXT_SEPARATOR = '\n\n'
const MAX_SUFFIX_RATIO = 0.35

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, FALLBACK_BLOCK_SEPARATOR, FALLBACK_LEAF_TEXT)
}

function getJoinedLength(parts: string[]) {
  let total = 0
  let hasContent = false
  for (const part of parts) {
    if (!part) continue
    if (hasContent) total += CONTEXT_SEPARATOR.length
    total += part.length
    hasContent = true
  }
  return total
}

function takeTail(text: string, limit: number) {
  if (!text || limit <= 0) return ''
  if (text.length <= limit) return text
  return text.slice(-limit)
}

function takeHead(text: string, limit: number) {
  if (!text || limit <= 0) return ''
  if (text.length <= limit) return text
  return text.slice(0, limit)
}

function fitPrefixSections(parts: string[], limit: number) {
  if (limit <= 0) return ''

  const fitted: string[] = []
  let remaining = limit

  for (let index = parts.length - 1; index >= 0; index -= 1) {
    const part = parts[index]
    if (!part || remaining <= 0) continue

    const separatorCost = fitted.length > 0 ? CONTEXT_SEPARATOR.length : 0
    if (remaining <= separatorCost) break

    const nextPart = takeTail(part, remaining - separatorCost)
    if (!nextPart) continue

    fitted.unshift(nextPart)
    remaining -= nextPart.length + separatorCost
  }

  return fitted.join(CONTEXT_SEPARATOR)
}

function fitSuffixSections(parts: string[], limit: number) {
  if (limit <= 0) return ''

  const fitted: string[] = []
  let remaining = limit

  for (const part of parts) {
    if (!part || remaining <= 0) continue

    const separatorCost = fitted.length > 0 ? CONTEXT_SEPARATOR.length : 0
    if (remaining <= separatorCost) break

    const nextPart = takeHead(part, remaining - separatorCost)
    if (!nextPart) continue

    fitted.push(nextPart)
    remaining -= nextPart.length + separatorCost
  }

  return fitted.join(CONTEXT_SEPARATOR)
}

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 prefixParts = [docContext, before, payload.prefix].filter(Boolean)
    const suffixParts = [payload.suffix, after].filter(Boolean)
    const mergedPrefix = prefixParts.join(CONTEXT_SEPARATOR)
    const mergedSuffix = suffixParts.join(CONTEXT_SEPARATOR)

    if (mergedPrefix.length + mergedSuffix.length > DOC_CONTEXT_LIMIT) {
      const prefixCapacity = getJoinedLength(prefixParts)
      const suffixCapacity = getJoinedLength(suffixParts)
      const maxSuffixBudget = Math.min(suffixCapacity, Math.floor(DOC_CONTEXT_LIMIT * MAX_SUFFIX_RATIO))

      let suffixBudget = maxSuffixBudget
      let prefixBudget = DOC_CONTEXT_LIMIT - suffixBudget

      if (prefixCapacity < prefixBudget) {
        const transferable = prefixBudget - prefixCapacity
        suffixBudget = Math.min(suffixCapacity, suffixBudget + transferable)
        prefixBudget = DOC_CONTEXT_LIMIT - suffixBudget
      } else if (suffixCapacity < suffixBudget) {
        const transferable = suffixBudget - suffixCapacity
        prefixBudget = Math.min(prefixCapacity, prefixBudget + transferable)
        suffixBudget = DOC_CONTEXT_LIMIT - prefixBudget
      }

      return {
        prefix: fitPrefixSections(prefixParts, prefixBudget),
        suffix: fitSuffixSections(suffixParts, suffixBudget),
        languageId: payload.languageId,
        blocked: false,
      }
    }

    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)
}