Atomizer/atomizer-dashboard/frontend/src/hooks/useSpecWebSocket.ts

/**
 * useSpecWebSocket - WebSocket connection for real-time spec sync
 *
 * Connects to the backend WebSocket endpoint for live spec updates.
 * Handles auto-reconnection, message parsing, and store updates.
 *
 * P2.11-P2.14: WebSocket sync implementation
 */

import { useEffect, useRef, useCallback, useState } from 'react';
import { useSpecStore } from './useSpecStore';

// ============================================================================
// Types
// ============================================================================

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

interface SpecWebSocketMessage {
  type: 'modification' | 'full_sync' | 'error' | 'ping';
  payload: unknown;
}

interface ModificationPayload {
  operation: 'set' | 'add' | 'remove';
  path: string;
  value?: unknown;
  modified_by: string;
  timestamp: string;
  hash: string;
}

interface ErrorPayload {
  message: string;
  code?: string;
}

interface UseSpecWebSocketOptions {
  /**
   * Enable auto-reconnect on disconnect (default: true)
   */
  autoReconnect?: boolean;

  /**
   * Reconnect delay in ms (default: 3000)
   */
  reconnectDelay?: number;

  /**
   * Max reconnect attempts (default: 10)
   */
  maxReconnectAttempts?: number;

  /**
   * Client identifier for tracking modifications (default: 'canvas')
   */
  clientId?: string;
}

interface UseSpecWebSocketReturn {
  /**
   * Current connection status
   */
  status: ConnectionStatus;

  /**
   * Manually disconnect
   */
  disconnect: () => void;

  /**
   * Manually reconnect
   */
  reconnect: () => void;

  /**
   * Send a message to the WebSocket (for future use)
   */
  send: (message: SpecWebSocketMessage) => void;

  /**
   * Last error message if any
   */
  lastError: string | null;
}

// ============================================================================
// Hook
// ============================================================================

export function useSpecWebSocket(
  studyId: string | null,
  options: UseSpecWebSocketOptions = {}
): UseSpecWebSocketReturn {
  const {
    autoReconnect = true,
    reconnectDelay = 3000,
    maxReconnectAttempts = 10,
    clientId = 'canvas',
  } = options;

  const wsRef = useRef<WebSocket | null>(null);
  const reconnectAttemptsRef = useRef(0);
  const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);

  const [status, setStatus] = useState<ConnectionStatus>('disconnected');
  const [lastError, setLastError] = useState<string | null>(null);

  // Get store actions
  const reloadSpec = useSpecStore((s) => s.reloadSpec);
  const setError = useSpecStore((s) => s.setError);

  // Build WebSocket URL
  const getWsUrl = useCallback((id: string): string => {
    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
    const host = window.location.host;
    return `${protocol}//${host}/api/studies/${encodeURIComponent(id)}/spec/sync?client_id=${clientId}`;
  }, [clientId]);

  // Handle incoming messages
  const handleMessage = useCallback((event: MessageEvent) => {
    try {
      const message: SpecWebSocketMessage = JSON.parse(event.data);

      switch (message.type) {
        case 'modification': {
          const payload = message.payload as ModificationPayload;

          // Skip if this is our own modification
          if (payload.modified_by === clientId) {
            return;
          }

          // Reload spec to get latest state
          // In a more sophisticated implementation, we could apply the patch locally
          reloadSpec().catch((err) => {
            console.error('Failed to reload spec after modification:', err);
          });
          break;
        }

        case 'full_sync': {
          // Full spec sync requested (e.g., after reconnect)
          reloadSpec().catch((err) => {
            console.error('Failed to reload spec during full_sync:', err);
          });
          break;
        }

        case 'error': {
          const payload = message.payload as ErrorPayload;
          console.error('WebSocket error:', payload.message);
          setLastError(payload.message);
          setError(payload.message);
          break;
        }

        case 'ping': {
          // Keep-alive ping, respond with pong
          if (wsRef.current?.readyState === WebSocket.OPEN) {
            wsRef.current.send(JSON.stringify({ type: 'pong' }));
          }
          break;
        }

        default:
          console.warn('Unknown WebSocket message type:', message.type);
      }
    } catch (error) {
      console.error('Failed to parse WebSocket message:', error);
    }
  }, [clientId, reloadSpec, setError]);

  // Connect to WebSocket
  const connect = useCallback(() => {
    if (!studyId) return;

    // Clean up existing connection
    if (wsRef.current) {
      wsRef.current.close();
    }

    setStatus('connecting');
    setLastError(null);

    const url = getWsUrl(studyId);
    const ws = new WebSocket(url);

    ws.onopen = () => {
      setStatus('connected');
      reconnectAttemptsRef.current = 0;
    };

    ws.onmessage = handleMessage;

    ws.onerror = (event) => {
      console.error('WebSocket error:', event);
      setLastError('WebSocket connection error');
    };

    ws.onclose = (_event) => {
      setStatus('disconnected');

      // Check if we should reconnect
      if (autoReconnect && reconnectAttemptsRef.current < maxReconnectAttempts) {
        reconnectAttemptsRef.current++;
        setStatus('reconnecting');

        // Clear any existing reconnect timeout
        if (reconnectTimeoutRef.current) {
          clearTimeout(reconnectTimeoutRef.current);
        }

        // Schedule reconnect with exponential backoff
        const delay = reconnectDelay * Math.min(reconnectAttemptsRef.current, 5);
        reconnectTimeoutRef.current = setTimeout(() => {
          connect();
        }, delay);
      } else if (reconnectAttemptsRef.current >= maxReconnectAttempts) {
        setLastError('Max reconnection attempts reached');
      }
    };

    wsRef.current = ws;
  }, [studyId, getWsUrl, handleMessage, autoReconnect, reconnectDelay, maxReconnectAttempts]);

  // Disconnect
  const disconnect = useCallback(() => {
    // Clear reconnect timeout
    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current);
      reconnectTimeoutRef.current = null;
    }

    // Close WebSocket
    if (wsRef.current) {
      wsRef.current.close();
      wsRef.current = null;
    }

    reconnectAttemptsRef.current = maxReconnectAttempts; // Prevent auto-reconnect
    setStatus('disconnected');
  }, [maxReconnectAttempts]);

  // Reconnect
  const reconnect = useCallback(() => {
    reconnectAttemptsRef.current = 0;
    connect();
  }, [connect]);

  // Send message
  const send = useCallback((message: SpecWebSocketMessage) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(message));
    } else {
      console.warn('WebSocket not connected, cannot send message');
    }
  }, []);

  // Connect when studyId changes
  useEffect(() => {
    if (studyId) {
      connect();
    } else {
      disconnect();
    }

    return () => {
      // Cleanup on unmount or studyId change
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
      }
      if (wsRef.current) {
        wsRef.current.close();
      }
    };
  }, [studyId, connect, disconnect]);

  return {
    status,
    disconnect,
    reconnect,
    send,
    lastError,
  };
}

export default useSpecWebSocket;