calctext/calcpad-web/src/collab/useWebSocketProvider.ts

import { useRef, useEffect, useState } from 'react'
import * as Y from 'yjs'
import { WebsocketProvider } from 'y-websocket'
import type { Awareness } from 'y-protocols/awareness'

const COLLAB_WS_URL = import.meta.env.VITE_COLLAB_WS_URL as string || 'ws://localhost:4000'

interface UseWebSocketProviderOptions {
  ydoc: Y.Doc | null
  roomName: string | null
  token: string | null
  userName: string
  userColor: string
}

/**
 * Manages a Y.js WebSocket provider for real-time collaboration.
 * Connects to the Hocuspocus server.
 */
export function useWebSocketProvider({
  ydoc,
  roomName,
  token,
  userName,
  userColor,
}: UseWebSocketProviderOptions) {
  const providerRef = useRef<WebsocketProvider | null>(null)
  const [awareness, setAwareness] = useState<Awareness | null>(null)
  const [connected, setConnected] = useState(false)
  const [peerCount, setPeerCount] = useState(0)

  useEffect(() => {
    if (!ydoc || !roomName || !token) {
      setAwareness(null)
      setConnected(false)
      setPeerCount(0)
      return
    }

    const provider = new WebsocketProvider(
      COLLAB_WS_URL,
      roomName,
      ydoc,
      {
        params: { token },
        connect: true,
      },
    )

    providerRef.current = provider

    // Set local user awareness
    provider.awareness.setLocalStateField('user', {
      name: userName,
      color: userColor,
    })

    setAwareness(provider.awareness)

    // Track connection status
    provider.on('status', (event: { status: string }) => {
      setConnected(event.status === 'connected')
    })

    // Track peer count
    const updatePeerCount = () => {
      const states = provider.awareness.getStates()
      setPeerCount(Math.max(0, states.size - 1)) // Exclude self
    }

    provider.awareness.on('change', updatePeerCount)
    updatePeerCount()

    return () => {
      provider.awareness.off('change', updatePeerCount)
      provider.disconnect()
      provider.destroy()
      providerRef.current = null
      setAwareness(null)
      setConnected(false)
      setPeerCount(0)
    }
  }, [ydoc, roomName, token, userName, userColor])

  return { awareness, connected, peerCount }
}