smoothschedule/frontend/src/hooks/useUserNotifications.ts

/**
 * WebSocket hook for real-time user notifications.
 * Connects to the backend WebSocket and updates React Query cache.
 */

import { useEffect, useRef, useCallback } from 'react';
import { useQueryClient } from '@tanstack/react-query';
import { getCookie } from '../utils/cookies';
import { getWebSocketUrl } from '../utils/domain';
import { UserEmail } from '../api/profile';

interface WebSocketMessage {
  type: 'connection_established' | 'email_verified' | 'profile_updated' | 'pong';
  email_id?: number;
  email?: string;
  user_id?: string;
  message?: string;
  fields?: string[];
}

interface UseUserNotificationsOptions {
  enabled?: boolean;
  onConnected?: () => void;
  onDisconnected?: () => void;
  onError?: (error: Event) => void;
  onEmailVerified?: (emailId: number, email: string) => void;
}

/**
 * Hook for real-time user notifications via WebSocket.
 */
export function useUserNotifications(options: UseUserNotificationsOptions = {}) {
  const { enabled = true, onConnected, onDisconnected, onError, onEmailVerified } = options;
  const queryClient = useQueryClient();

  // Use refs for callbacks to avoid recreating connect function
  const onConnectedRef = useRef(onConnected);
  const onDisconnectedRef = useRef(onDisconnected);
  const onErrorRef = useRef(onError);
  const onEmailVerifiedRef = useRef(onEmailVerified);

  // Update refs when callbacks change
  useEffect(() => {
    onConnectedRef.current = onConnected;
    onDisconnectedRef.current = onDisconnected;
    onErrorRef.current = onError;
    onEmailVerifiedRef.current = onEmailVerified;
  }, [onConnected, onDisconnected, onError, onEmailVerified]);

  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
  const pingIntervalRef = useRef<NodeJS.Timeout | null>(null);
  const reconnectAttemptsRef = useRef(0);
  const isConnectingRef = useRef(false);
  const maxReconnectAttempts = 5;

  const updateEmailQueryCache = useCallback((emailId: number) => {
    // Update the userEmails query cache to mark the email as verified
    queryClient.setQueryData<UserEmail[]>(['userEmails'], (old) => {
      if (!old) return old;
      return old.map((email) =>
        email.id === emailId ? { ...email, verified: true } : email
      );
    });
  }, [queryClient]);

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

    // Clear ping interval
    if (pingIntervalRef.current) {
      clearInterval(pingIntervalRef.current);
      pingIntervalRef.current = null;
    }

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

    isConnectingRef.current = false;
  }, []);

  const connect = useCallback(() => {
    // Prevent multiple simultaneous connection attempts
    if (isConnectingRef.current || wsRef.current?.readyState === WebSocket.OPEN) {
      return;
    }

    const token = getCookie('access_token');
    if (!token) {
      console.log('UserNotifications WebSocket: Missing token, skipping connection');
      return;
    }

    isConnectingRef.current = true;

    // Close existing connection if any
    if (wsRef.current) {
      wsRef.current.close();
    }

    // Build WebSocket URL dynamically
    const url = getWebSocketUrl(`user/?token=${token}`);

    console.log('UserNotifications WebSocket: Connecting');
    const ws = new WebSocket(url);

    ws.onopen = () => {
      console.log('UserNotifications WebSocket: Connected');
      reconnectAttemptsRef.current = 0;
      isConnectingRef.current = false;
      onConnectedRef.current?.();

      // Start ping interval to keep connection alive
      pingIntervalRef.current = setInterval(() => {
        if (ws.readyState === WebSocket.OPEN) {
          ws.send(JSON.stringify({ type: 'ping' }));
        }
      }, 30000); // Ping every 30 seconds
    };

    ws.onmessage = (event) => {
      try {
        const message: WebSocketMessage = JSON.parse(event.data);

        switch (message.type) {
          case 'connection_established':
            console.log('UserNotifications WebSocket: Connection confirmed -', message.message);
            break;
          case 'pong':
            // Heartbeat response, ignore
            break;
          case 'email_verified':
            console.log('UserNotifications WebSocket: Email verified', message.email);
            if (message.email_id) {
              updateEmailQueryCache(message.email_id);
              onEmailVerifiedRef.current?.(message.email_id, message.email || '');
            }
            break;
          case 'profile_updated':
            console.log('UserNotifications WebSocket: Profile updated', message.fields);
            // Invalidate profile queries to refresh data
            queryClient.invalidateQueries({ queryKey: ['currentUser'] });
            break;
          default:
            console.log('UserNotifications WebSocket: Unknown message type', message);
        }
      } catch (err) {
        console.error('UserNotifications WebSocket: Failed to parse message', err);
      }
    };

    ws.onerror = (error) => {
      console.error('UserNotifications WebSocket: Error', error);
      isConnectingRef.current = false;
      onErrorRef.current?.(error);
    };

    ws.onclose = (event) => {
      console.log('UserNotifications WebSocket: Disconnected', event.code, event.reason);
      isConnectingRef.current = false;
      onDisconnectedRef.current?.();

      // Clear ping interval
      if (pingIntervalRef.current) {
        clearInterval(pingIntervalRef.current);
        pingIntervalRef.current = null;
      }

      // Only attempt reconnection if this wasn't a deliberate close
      // Code 1000 = normal closure, 1001 = going away (page unload)
      if (event.code !== 1000 && event.code !== 1001 && reconnectAttemptsRef.current < maxReconnectAttempts) {
        const delay = Math.min(1000 * Math.pow(2, reconnectAttemptsRef.current), 30000);
        console.log(`UserNotifications WebSocket: Reconnecting in ${delay}ms (attempt ${reconnectAttemptsRef.current + 1})`);

        reconnectTimeoutRef.current = setTimeout(() => {
          reconnectAttemptsRef.current++;
          connect();
        }, delay);
      }
    };

    wsRef.current = ws;
  }, [queryClient, updateEmailQueryCache]);

  useEffect(() => {
    if (enabled) {
      connect();
    }

    return () => {
      disconnect();
    };
  }, [enabled, connect, disconnect]);

  return {
    isConnected: wsRef.current?.readyState === WebSocket.OPEN,
    reconnect: connect,
    disconnect,
  };
}