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.
 */
import * as Y from 'yjs';
import type { YjsDocument } from './YjsDocument';

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

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

export class WebSocketProvider {
  private ws: WebSocket | null = null;
  private url: string;
  private docName: string;
  private yjsDoc: YjsDocument;
  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;

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

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

    this.shouldReconnect = true;
    this.setStatus('connecting');

    const fullUrl = `${this.url}/ws/collab/${encodeURIComponent(this.docName)}`;
    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;
      // Send local state to server for initial sync
      const stateUpdate = Y.encodeStateAsUpdate(this.yjsDoc.doc);
      if (stateUpdate.length > 0) {
        this.ws?.send(stateUpdate);
      }
      this.onSync?.();
    };

    this.ws.onmessage = (event: MessageEvent) => {
      try {
        const data = new Uint8Array(event.data as ArrayBuffer);
        this.yjsDoc.applyUpdate(data);
      } catch (err) {
        console.warn('[WebSocketProvider] Failed to apply update:', err);
      }
    };

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

    this.ws.onclose = () => {
      this.ws = null;
      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.ws.send(update);
    }
  }

  /** 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);
    };
  }

  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.setStatus('disconnected');
  }

  private scheduleReconnect(): void {
    if (this.reconnectTimer) return;
    this.reconnectTimer = setTimeout(() => {
      this.reconnectTimer = null;
      this.connect();
    }, 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);
    }
  }
}