OJ2/apps/web/src/shared/composables/collabDoc.ts

import { Compartment } from '@codemirror/state'
import type { EditorView } from '@codemirror/view'
import { useCollabStore } from 'shared/store/collab'
import { useUserStore } from 'shared/store/user'

/** y-websocket 那套消息头，服务端不解析，只有两端认 */
const MESSAGE_SYNC = 0
const MESSAGE_AWARENESS = 1

const TEACHER_COLOR = '#ff6b6b'
const STUDENT_COLOR = '#4dabf7'

interface StartOptions {
  editorView: EditorView
  /**
   * 文档的初始内容。
   *
   * **学生端传当前编辑器内容，教师端必须传 null。** 这是硬规则：
   * 求助是学生发起的，学生的代码是唯一内容源。老师端插入任何初始内容都会
   * 与学生的内容合并，结果就是两份代码拼在一起 —— 老实现的竞态就是这么来的。
   */
  seedContent: string | null
}

export function useCollabDoc() {
  const collabStore = useCollabStore()
  const userStore = useUserStore()
  const compartment = new Compartment()

  let doc: any = null
  let awareness: any = null
  let view: EditorView | null = null
  let detachDocUpdate: (() => void) | null = null
  let detachAwarenessUpdate: (() => void) | null = null
  /**
   * 会话代号。每次 start / stop 都推进一格。
   *
   * start 要 await 六个动态 import，房间可能在这期间就关了（机房首次加载
   * y* 那几个 chunk 正是最慢的时候）。stop() 先跑完的话它面对的是 doc === null，
   * 什么也拆不到；等 import 回来 start 的后半段照样建文档、装 binaryHandler、
   * 往编辑器里挂 yCollab、还发一轮 SyncStep1 —— 给一个已经不存在的房间。
   */
  let generation = 0

  async function start({ editorView, seedContent }: StartOptions) {
    // 已经有一份在跑就先收掉。重复 start 只可能来自时序竞争，叠加没有意义
    if (doc) stop()
    const myGeneration = ++generation

    const [Y, awarenessProtocol, syncProtocol, encoding, decoding, { yCollab }] = await Promise.all([
      import('yjs'),
      import('y-protocols/awareness'),
      import('y-protocols/sync'),
      import('lib0/encoding'),
      import('lib0/decoding'),
      import('y-codemirror.next'),
    ])

    // 等 chunk 的这段时间里房间关了（或者又开了新的一轮），整个放弃
    if (myGeneration !== generation) return

    view = editorView
    doc = new Y.Doc()
    const ytext = doc.getText('codemirror')
    awareness = new awarenessProtocol.Awareness(doc)

    // ★ 先把内容写进 ytext，再挂 yCollab。yCollab 挂上去**不会**拿 ytext 覆盖编辑器，
    // 只同步此后的改动 —— 所以两边挂之前编辑器里的内容必须和 ytext 一致：
    // 学生端 ytext = 编辑器内容（种子）；教师端 ytext 是空的，编辑器也得清空。
    if (seedContent) ytext.insert(0, seedContent)
    // 教师端编辑器里原本是老师自己这道题的草稿或模板。不清的话它一直留在老师那边，学生的代码
    // 同步过来插在它前面，两边从此不一致（实测：老师看到两行一样的 input()，学生只有一行）。
    // 清空会顺着 v-model 写回 codeStore，但协作中的老师不存草稿（ProblemEditor 的 persistCode），
    // 结束后 loadCode 会把他的草稿读回来
    if (seedContent === null && editorView.state.doc.length) {
      editorView.dispatch({ changes: { from: 0, to: editorView.state.doc.length, insert: '' } })
    }

    const send = (build: (encoder: any) => void) => {
      const encoder = encoding.createEncoder()
      build(encoder)
      collabStore.sendBinary(encoding.toUint8Array(encoder))
    }

    collabStore.setBinaryHandler((data) => {
      const decoder = decoding.createDecoder(new Uint8Array(data))
      const messageType = decoding.readVarUint(decoder)
      if (messageType === MESSAGE_SYNC) {
        const encoder = encoding.createEncoder()
        encoding.writeVarUint(encoder, MESSAGE_SYNC)
        syncProtocol.readSyncMessage(decoder, encoder, doc, 'remote')
        // 只有需要回话时才发（readSyncMessage 可能什么都没写）
        if (encoding.length(encoder) > 1) {
          collabStore.sendBinary(encoding.toUint8Array(encoder))
        }
      } else if (messageType === MESSAGE_AWARENESS) {
        awarenessProtocol.applyAwarenessUpdate(awareness, decoding.readVarUint8Array(decoder), 'remote')
      }
    })

    const onDocUpdate = (update: Uint8Array, origin: any) => {
      if (origin === 'remote') return
      send((encoder) => {
        encoding.writeVarUint(encoder, MESSAGE_SYNC)
        syncProtocol.writeUpdate(encoder, update)
      })
    }
    doc.on('update', onDocUpdate)
    detachDocUpdate = () => doc?.off('update', onDocUpdate)

    const onAwarenessUpdate = ({ added, updated, removed }: { added: number[]; updated: number[]; removed: number[] }, origin: any) => {
      if (origin === 'remote') return
      const changed = added.concat(updated, removed)
      send((encoder) => {
        encoding.writeVarUint(encoder, MESSAGE_AWARENESS)
        encoding.writeVarUint8Array(encoder, awarenessProtocol.encodeAwarenessUpdate(awareness, changed))
      })
    }
    awareness.on('update', onAwarenessUpdate)
    detachAwarenessUpdate = () => awareness?.off('update', onAwarenessUpdate)

    awareness.setLocalStateField('user', {
      name: userStore.user?.username ?? '匿名',
      color: userStore.isTeacherOrAbove ? TEACHER_COLOR : STUDENT_COLOR,
    })

    editorView.dispatch({
      effects: compartment.reconfigure(yCollab(ytext, awareness)),
    })

    // 握手：双方都发 SyncStep1，各自回 Step2，两边收敛。
    // 服务端是哑转发，不参与同步，所以这一步必须由两端对称完成
    send((encoder) => {
      encoding.writeVarUint(encoder, MESSAGE_SYNC)
      syncProtocol.writeSyncStep1(encoder, doc)
    })
    send((encoder) => {
      encoding.writeVarUint(encoder, MESSAGE_AWARENESS)
      encoding.writeVarUint8Array(encoder, awarenessProtocol.encodeAwarenessUpdate(awareness, [doc.clientID]))
    })
  }

  function stop() {
    generation += 1
    collabStore.setBinaryHandler(null)
    detachDocUpdate?.()
    detachAwarenessUpdate?.()
    detachDocUpdate = null
    detachAwarenessUpdate = null

    if (view) {
      try {
        view.dispatch({ effects: compartment.reconfigure([]) })
      } catch (error) {
        console.warn('移除协同编辑扩展失败:', error)
      }
      view = null
    }
    awareness?.destroy()
    doc?.destroy()
    awareness = null
    doc = null
  }

  function getInitialExtension() {
    return compartment.of([])
  }

  return { start, stop, getInitialExtension }
}