web-cad/frontend/src/crdt/WebSocketProvider.ts

/**
 * WebSocketProvider – Custom WebSocket provider matching the backend raw-update protocol.
 * Backend sends Y.encodeStateAsUpdate on connect and applies raw updates on message.
 *
 * Issue #20: On connect, the client must NOT send its local state immediately.
 * Instead, the server sends its state first (Y.encodeStateAsUpdate). The client
 * applies that update to sync, then sends only incremental local updates.
 */
import * as Y from 'yjs';
import type { YjsDocument } from './YjsDocument';
import type { AwarenessManager } from './AwarenessManager';

export type ConnectionStatus = 'disconnected' | 'connecting' | 'connected' | 'error';

export interface WebSocketProviderOptions {
  url: string;
  docName: string;
  yjsDoc: YjsDocument;
  awareness?: AwarenessManager;
  onStatusChange?: (status: ConnectionStatus) => void;
  onSync?: () => void;
}

// Message type prefixes for multiplexing Y.Doc updates and awareness updates
// Using a simple byte-prefix scheme: 0 = Y.Doc update, 1 = awareness update
const MSG_DOC_UPDATE = 0;
const MSG_AWARENESS = 1;

export class WebSocketProvider {
  private ws: WebSocket | null = null;
  private url: string;
  private docName: string;
  private yjsDoc: YjsDocument;
  private awareness?: AwarenessManager;
  private status: ConnectionStatus = 'disconnected';
  private onStatusChange?: (status: ConnectionStatus) => void;
  private onSync?: () => void;
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
  private reconnectDelay = 1000;
  private maxReconnectDelay = 30000;
  private shouldReconnect = true;
  private authToken: string | undefined;
  private synced = false;

  constructor(opts: WebSocketProviderOptions) {
    this.url = opts.url;
    this.docName = opts.docName;
    this.yjsDoc = opts.yjsDoc;
    this.awareness = opts.awareness;
    this.onStatusChange = opts.onStatusChange;
    this.onSync = opts.onSync;
  }

  connect(token?: string): void {
    if (this.ws && (this.ws.readyState === WebSocket.OPEN || this.ws.readyState === WebSocket.CONNECTING)) {
      return;
    }

    this.authToken = token;
    this.shouldReconnect = true;
    this.synced = false;
    this.setStatus('connecting');

    const tokenParam = token ? `?token=${encodeURIComponent(token)}` : '';
    const fullUrl = `${this.url}/ws/collab/${encodeURIComponent(this.docName)}${tokenParam}`;
    try {
      this.ws = new WebSocket(fullUrl);
    } catch (err) {
      console.warn('[WebSocketProvider] Failed to construct WebSocket:', err);
      this.setStatus('error');
      this.scheduleReconnect();
      return;
    }
    this.ws.binaryType = 'arraybuffer';

    this.ws.onopen = () => {
      this.setStatus('connected');
      this.reconnectDelay = 1000;
      // Issue #20: Do NOT send local state on connect.
      // The server sends its state (Y.encodeStateAsUpdate) to new clients.
      // We wait for that server state, apply it, then mark as synced.
      // Only after sync do we forward incremental local updates.
      //
      // If we have awareness, send our local awareness state after connect.
      if (this.awareness) {
        const awarenessUpdate = this.awareness.encodeLocalState();
        if (awarenessUpdate && awarenessUpdate.length > 0) {
          this.sendAwarenessUpdate(awarenessUpdate);
        }
      }
    };

    this.ws.onmessage = (event: MessageEvent) => {
      try {
        const data = new Uint8Array(event.data as ArrayBuffer);
        if (data.length === 0) return;

        // Check message type prefix (first byte)
        const msgType = data[0];
        const payload = data.slice(1);

        if (msgType === MSG_AWARENESS) {
          // Awareness update — apply to awareness manager
          if (this.awareness) {
            this.awareness.applyRemoteUpdate(payload, 'remote');
          }
        } else {
          // Y.Doc update — apply to Y.Doc
          // Support both prefixed and non-prefixed messages for backward compatibility
          // If the first byte looks like a valid Y.Doc update, use the full data;
          // otherwise treat the whole message as the Y.Doc update (legacy mode).
          const updateData = msgType === MSG_DOC_UPDATE ? payload : data;
          this.yjsDoc.applyUpdate(updateData);

          // Mark as synced after receiving the first server state
          if (!this.synced) {
            this.synced = true;
            this.onSync?.();
          }
        }
      } catch (err) {
        console.warn('[WebSocketProvider] Failed to process message:', err);
      }
    };

    this.ws.onerror = () => {
      this.setStatus('error');
    };

    this.ws.onclose = () => {
      this.ws = null;
      this.synced = false;
      this.setStatus('disconnected');
      if (this.shouldReconnect) {
        this.scheduleReconnect();
      }
    };
  }

  /** Send a local Y.Doc update to the server */
  sendUpdate(update: Uint8Array): void {
    if (this.ws && this.ws.readyState === WebSocket.OPEN && this.synced) {
      // Prefix with MSG_DOC_UPDATE byte
      const msg = new Uint8Array(update.length + 1);
      msg[0] = MSG_DOC_UPDATE;
      msg.set(update, 1);
      this.ws.send(msg);
    }
  }

  /** Send a local awareness update to the server */
  sendAwarenessUpdate(update: Uint8Array): void {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      const msg = new Uint8Array(update.length + 1);
      msg[0] = MSG_AWARENESS;
      msg.set(update, 1);
      this.ws.send(msg);
    }
  }

  /** Listen for local doc changes and forward to server */
  bindLocalUpdates(): () => void {
    const handler = (update: Uint8Array, origin: unknown) => {
      // Only send updates that originated locally (not from remote apply)
      if (origin !== 'remote') {
        this.sendUpdate(update);
      }
    };
    this.yjsDoc.doc.on('update', handler);
    return () => {
      this.yjsDoc.doc.off('update', handler);
    };
  }

  /** Listen for local awareness changes and forward to server */
  bindAwarenessUpdates(): () => void {
    if (!this.awareness) return () => {};
    const handler = () => {
      const update = this.awareness!.encodeLocalState();
      if (update && update.length > 0) {
        this.sendAwarenessUpdate(update);
      }
    };
    this.awareness.awareness.on('update', handler);
    return () => {
      this.awareness!.awareness.off('update', handler);
    };
  }

  isSynced(): boolean {
    return this.synced;
  }

  getStatus(): ConnectionStatus {
    return this.status;
  }

  disconnect(): void {
    this.shouldReconnect = false;
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = null;
    }
    if (this.ws) {
      this.ws.close();
      this.ws = null;
    }
    this.synced = false;
    this.setStatus('disconnected');
  }

  private scheduleReconnect(): void {
    if (this.reconnectTimer) return;
    this.reconnectTimer = setTimeout(() => {
      this.reconnectTimer = null;
      this.connect(this.authToken);
    }, this.reconnectDelay);
    this.reconnectDelay = Math.min(this.reconnectDelay * 2, this.maxReconnectDelay);
  }

  private setStatus(status: ConnectionStatus): void {
    if (this.status !== status) {
      this.status = status;
      this.onStatusChange?.(status);
    }
  }
}