yjs-bpmn-frontend/src/composables/useCollaborativeFlow.ts

import type { Edge, Node } from '@vue-flow/core'
import type { MaybeRefOrGetter } from 'vue'
import type { YMapEvent } from 'yjs'
import { useVueFlow } from '@vue-flow/core'
import { onScopeDispose, shallowRef, toValue, watch } from 'vue'
import { WebsocketProvider } from 'y-websocket'
import * as Y from 'yjs'

export function useCollaborativeFlow(diagramId: MaybeRefOrGetter<string>) {
  const yDoc = new Y.Doc()

  const provider = shallowRef<WebsocketProvider>()

  const yNodes = yDoc.getMap<Node>('nodes')
  const yEdges = yDoc.getMap<Edge>('edges')

  const {
    addNodes,
    addEdges,
    updateNode,
    onNodesChange,
    onEdgesChange,
    applyNodeChanges,
    applyEdgeChanges,
    removeNodes,
    removeEdges,
  } = useVueFlow()

  let isApplyingFromY = false
  let isApplyingEdgesFromY = false

  // ---------------------------
  // Yjs -> Vue Flow
  // ---------------------------
  const nodesObserver = (event: YMapEvent<Node>) => {
    console.log('nodesObserver', event)

    isApplyingFromY = true

    for (const [id, change] of event.changes.keys) {
      switch (change.action) {
        case 'add': {
          const node = yNodes.get(id)

          if (!node)
            continue

          addNodes(node)
          break
        }

        case 'delete':
          removeNodes(id)
          break

        case 'update': {
          const node = yNodes.get(id)

          if (!node)
            continue

          updateNode(id, { position: node.position })
          break
        }
      }
    }

    isApplyingFromY = false
  }

  const edgesObserver = (event: YMapEvent<Edge>) => {
    console.log('edgesObserver', event)

    isApplyingEdgesFromY = true

    for (const [id, change] of event.changes.keys) {
      switch (change.action) {
        case 'add': {
          const edge = yEdges.get(id)

          if (!edge)
            continue

          addEdges(edge)
          break
        }

        case 'delete':
          removeEdges(id)
          break

        // case 'update':
        //   updateEdgeData(id, data)
        //   break
      }
    }

    isApplyingEdgesFromY = false
  }

  // ---------------------------
  // Vue Flow -> Yjs
  // ---------------------------
  onNodesChange((changes) => {
    console.log('onNodesChange', changes)

    if (changes.length === 0)
      return

    applyNodeChanges(changes)

    if (isApplyingFromY)
      return

    yDoc.transact(() => {
      for (const change of changes) {
        const { type } = change

        switch (type) {
          case 'add': {
            yNodes.set(change.item.id, change.item)
            break
          }

          case 'remove':
            yNodes.delete(change.id)
            break

          case 'position': {
            const node = yNodes.get(change.id)

            if (!node)
              return

            if (change.position) {
              node.position = change.position
            }

            yNodes.set(change.id, node)

            break
          }
        }
      }
    })
  })

  onEdgesChange((changes) => {
    console.log('onEdgesChange', changes)

    if (changes.length === 0)
      return

    applyEdgeChanges(changes)

    if (isApplyingEdgesFromY)
      return

    yDoc.transact(() => {
      for (const change of changes) {
        const { type } = change

        switch (type) {
          case 'add': {
            yEdges.set(change.item.id, change.item)
            break
          }

          case 'remove':
            yEdges.delete(change.id)
            break
        }
      }
    })
  })

  yNodes.observe(nodesObserver)
  yEdges.observe(edgesObserver)

  watch(() => toValue(diagramId), (diagramId) => {
    provider.value?.destroy()

    reset()

    if (!diagramId)
      return

    provider.value = new WebsocketProvider(
      import.meta.env.DEV ? 'ws://localhost:1234' : 'wss://api/koptilnya.xyz/bpmn',
      diagramId,
      yDoc,
    )
  }, { immediate: true })

  onScopeDispose(() => {
    yNodes.unobserve(nodesObserver)
    yEdges.unobserve(edgesObserver)

    provider.value?.destroy()
  })

  function reset() {
    // yDoc.destroy()
    yNodes.clear()
    yEdges.clear()
  }

  return {
    reset,
  }
}