PPT/src/views/Student/CustomWebsocketProvider.ts

import type * as Y from 'yjs'
import * as syncProtocol from 'y-protocols/sync'
import * as awarenessProtocol from 'y-protocols/awareness'
import * as encoding from 'lib0/encoding'
import * as decoding from 'lib0/decoding'
import { Observable } from 'lib0/observable'

const messageSync = 0
const messageAwareness = 1
const messageAuth = 2
const messageQueryAwareness = 3

export interface YjsRedisProviderOptions {
  params?: Record<string, string>
  awareness?: awarenessProtocol.Awareness
  resyncInterval?: number
  maxBackoffTime?: number
  connect?: boolean
}

/**
 * 兼容 y-websocket WebsocketProvider 接口的自定义实现。
 * 关键改进：
 *  1. readMessage 循环解析整条 buffer，兼容多消息合并推送
 *  2. 额外 emit 'raw-message'，把完整 Uint8Array 抛出，供上层解析
 *  3. 保留 ws / awareness / params / synced / connect / destroy 等原 API
 */
export class YjsRedisProvider extends Observable<string> {
  public doc: Y.Doc
  public awareness: awarenessProtocol.Awareness
  public ws: WebSocket | null = null
  public wsconnected = false
  public wsconnecting = false
  public synced = false
  public shouldConnect = true
  public params: Record<string, string>

  private serverUrl: string
  private roomname: string
  private url: string
  private resyncInterval: number
  private maxBackoffTime: number
  private reconnectAttempts = 0
  private _resyncIntervalId: any = null
  private _updateHandler: (update: Uint8Array, origin: any) => void
  private _awarenessUpdateHandler: (...args: any[]) => void
  private messageHandlers: Array<(
    encoder: encoding.Encoder,
    decoder: decoding.Decoder,
    provider: YjsRedisProvider,
    emitSynced: boolean,
    messageType: number
  ) => void>

  constructor(serverUrl: string, roomname: string, doc: Y.Doc, opts: YjsRedisProviderOptions = {}) {
    super()
    this.doc = doc
    this.roomname = roomname
    this.serverUrl = serverUrl.replace(/\/+$/, '')
    this.params = opts.params || {}
    this.resyncInterval = opts.resyncInterval ?? -1
    this.maxBackoffTime = opts.maxBackoffTime ?? 2500
    this.awareness = opts.awareness || new awarenessProtocol.Awareness(doc)
    this.awareness.setLocalState({})
    this.url = this._buildUrl()

    this.messageHandlers = []

    // 0: Sync
    this.messageHandlers[messageSync] = (encoder, decoder, provider, emitSynced) => {
      encoding.writeVarUint(encoder, messageSync)
      const syncMessageType = syncProtocol.readSyncMessage(decoder, encoder, doc, provider)
      if (emitSynced && syncMessageType === syncProtocol.messageYjsSyncStep2 && !provider.synced) {
        provider.synced = true
        provider.emit('sync', [true])
      }
    }

    // 1: Awareness
    this.messageHandlers[messageAwareness] = (encoder, decoder, provider) => {
      try {
        awarenessProtocol.applyAwarenessUpdate(
          provider.awareness,
          decoding.readVarUint8Array(decoder),
          provider
        )
      }
      catch (e) {
        // awareness 可能已被上层 destroy，忽略
      }
    }

    // 2: Auth（y-redis 返回权限信息）
    this.messageHandlers[messageAuth] = (encoder, decoder, provider) => {
      try {
        const authType = decoding.readVarUint(decoder)
        if (authType === 0) {
          provider.emit('auth-message', ['permission-denied'])
        }
        else if (authType === 1) {
          provider.emit('auth-message', ['authenticated'])
        }
      }
      catch (e) { /* ignore */ }
    }

    // 3: QueryAwareness
    this.messageHandlers[messageQueryAwareness] = (encoder, decoder, provider) => {
      encoding.writeVarUint(encoder, messageAwareness)
      encoding.writeVarUint8Array(
        encoder,
        awarenessProtocol.encodeAwarenessUpdate(
          provider.awareness,
          Array.from(provider.awareness.getStates().keys())
        )
      )
    }

    // 本地文档更新 → 广播
    this._updateHandler = (update, origin) => {
      if (origin !== this) {
        const encoder = encoding.createEncoder()
        encoding.writeVarUint(encoder, messageSync)
        syncProtocol.writeUpdate(encoder, update)
        this._broadcast(encoding.toUint8Array(encoder))
      }
    }
    this.doc.on('update', this._updateHandler)

    // 本地 awareness 更新 → 广播
    this._awarenessUpdateHandler = ({ added, updated, removed }, origin) => {
      const changedClients = added.concat(updated, removed)
      const encoder = encoding.createEncoder()
      encoding.writeVarUint(encoder, messageAwareness)
      encoding.writeVarUint8Array(
        encoder,
        awarenessProtocol.encodeAwarenessUpdate(this.awareness, changedClients)
      )
      this._broadcast(encoding.toUint8Array(encoder))
    }
    this.awareness.on('update', this._awarenessUpdateHandler)

    if (opts.connect !== false) {
      this.connect()
    }
  }

  private _buildUrl(): string {
    const qs = new URLSearchParams(this.params).toString()
    return `${this.serverUrl}/${this.roomname}${qs ? '?' + qs : ''}`
  }

  /** 更新参数（例如 token 过期后刷新 yauth），下次 connect 生效 */
  public updateParams(params: Record<string, string>): void {
    this.params = params
    this.url = this._buildUrl()
    this.emit('url-updated', [this.url])
  }

  public connect(): void {
    if (this.wsconnected || this.wsconnecting || !this.shouldConnect) return
    this.wsconnecting = true
    this.emit('status', [{ status: 'connecting' }])

    const ws = new WebSocket(this.url)
    this.ws = ws
    ws.binaryType = 'arraybuffer'

    ws.onopen = () => {
      this.wsconnecting = false
      this.wsconnected = true
      this.reconnectAttempts = 0
      this.emit('status', [{ status: 'connected' }])

      // Sync step 1
      const encoder = encoding.createEncoder()
      encoding.writeVarUint(encoder, messageSync)
      syncProtocol.writeSyncStep1(encoder, this.doc)
      ws.send(encoding.toUint8Array(encoder))

      // Awareness
      const states = this.awareness.getStates()
      if (states.size > 0) {
        const enc = encoding.createEncoder()
        encoding.writeVarUint(enc, messageAwareness)
        encoding.writeVarUint8Array(
          enc,
          awarenessProtocol.encodeAwarenessUpdate(this.awareness, Array.from(states.keys()))
        )
        ws.send(encoding.toUint8Array(enc))
      }

      // 启动周期 resync（如果用户传了 resyncInterval）
      if (this.resyncInterval > 0) {
        this._resyncIntervalId = setInterval(() => {
          if (this.ws && this.ws.readyState === WebSocket.OPEN) {
            const e = encoding.createEncoder()
            encoding.writeVarUint(e, messageSync)
            syncProtocol.writeSyncStep1(e, this.doc)
            this.ws.send(encoding.toUint8Array(e))
          }
        }, this.resyncInterval)
      }
    }

    ws.onmessage = (event: MessageEvent) => {
      if (event.data instanceof ArrayBuffer) {
        this._readMessage(new Uint8Array(event.data))
      }
      else if (event.data instanceof Blob) {
        event.data.arrayBuffer().then(buf => this._readMessage(new Uint8Array(buf)))
      }
      else if (typeof event.data === 'string') {
        this._readMessage(new TextEncoder().encode(event.data))
      }
    }

    ws.onclose = (event: CloseEvent) => {
      this.wsconnecting = false
      if (this._resyncIntervalId) {
        clearInterval(this._resyncIntervalId)
        this._resyncIntervalId = null
      }
      if (this.wsconnected) {
        this.wsconnected = false
        this.synced = false
        try {
          awarenessProtocol.removeAwarenessStates(
            this.awareness,
            Array.from(this.awareness.getStates().keys()).filter(
              client => client !== this.doc.clientID
            ),
            this
          )
        }
        catch (e) { /* ignore */ }
        this.emit('status', [{ status: 'disconnected' }])
      }
      this.emit('connection-close', [event, this])
      this.ws = null
      if (this.shouldConnect) {
        const delay = Math.min(this.maxBackoffTime, 100 * Math.pow(2, this.reconnectAttempts))
        this.reconnectAttempts++
        setTimeout(() => this.connect(), delay)
      }
    }

    ws.onerror = (event: Event) => {
      this.emit('connection-error', [event, this])
    }
  }

  /**
   * 核心：循环解析整条 buffer，兼容多消息合并
   * 额外：先抛出 raw-message，让上层拿到完整原始数据
   */
  private _readMessage(data: Uint8Array): void {
    // ★ 关键：把完整原始数据先广播出去
    this.emit('raw-message', [data])

    let offset = 0
    while (offset < data.length) {
      const decoder = decoding.createDecoder(data.subarray(offset))
      try {
        const messageType = decoding.readVarUint(decoder)
        const encoder = encoding.createEncoder()
        const handler = this.messageHandlers[messageType]

        if (handler) {
          handler(encoder, decoder, this, true, messageType)
          if (encoding.length(encoder) > 0 && this.ws?.readyState === WebSocket.OPEN) {
            this.ws.send(encoding.toUint8Array(encoder))
          }
        }

        const consumed = decoder.pos
        if (consumed === 0) break
        offset += consumed

        // 保留原有 message 事件，兼容旧监听（data 是这一条的片段）
        this.emit('message', [{
          data: data.slice(offset - consumed, offset),
          type: messageType,
        }])
      }
      catch (err) {
        // 解析失败立即停止，避免错位
        console.warn('[YjsRedisProvider] parse error at offset', offset, err)
        break
      }
    }
  }

  private _broadcast(data: Uint8Array): void {
    if (this.ws?.readyState === WebSocket.OPEN) {
      this.ws.send(data)
    }
  }

  public disconnect(): void {
    this.shouldConnect = false
    if (this._resyncIntervalId) {
      clearInterval(this._resyncIntervalId)
      this._resyncIntervalId = null
    }
    if (this.ws !== null) {
      this.ws.close()
    }
  }

  public destroy(): void {
    this.disconnect()
    this.doc.off('update', this._updateHandler)
    try {
      this.awareness.off('update', this._awarenessUpdateHandler) 
    }
    catch (e) { /* ignore */ }
    try {
      this.awareness.destroy() 
    }
    catch (e) { /* ignore */ }
    super.destroy()
  }
}