smoothschedule/frontend/src/hooks/useStaffEmailWebSocket.ts

import { useEffect, useRef, useCallback, useState } from 'react';
import { useQueryClient } from '@tanstack/react-query';
import { toast } from 'react-hot-toast';
import { useCurrentUser } from './useAuth';
import { getBaseDomain } from '../utils/domain';
import { getCookie } from '../utils/cookies';

/**
 * Event types sent by the staff email WebSocket consumer
 */
type StaffEmailEventType =
  | 'new_email'
  | 'email_read'
  | 'email_unread'
  | 'email_moved'
  | 'email_deleted'
  | 'folder_counts'
  | 'sync_started'
  | 'sync_completed'
  | 'sync_error';

interface NewEmailData {
  id?: number;
  subject?: string;
  from_name?: string;
  from_address?: string;
  snippet?: string;
  folder_id?: number;
  email_address_id?: number;
}

interface FolderCountData {
  [folderId: string]: {
    unread_count?: number;
    total_count?: number;
    folder_type?: string;
  };
}

interface SyncStatusData {
  email_address_id?: number;
  results?: Record<string, number>;
  new_count?: number;
  message?: string;
  details?: {
    results?: Record<string, number>;
    new_count?: number;
    message?: string;
  };
}

type StaffEmailData = NewEmailData | FolderCountData | SyncStatusData;

interface StaffEmailWebSocketMessage {
  type: StaffEmailEventType;
  data: StaffEmailData;
}

interface UseStaffEmailWebSocketOptions {
  /** Show toast notifications for events (default: true) */
  showToasts?: boolean;
  /** Callback when a new email arrives */
  onNewEmail?: (data: NewEmailData) => void;
  /** Callback when sync completes */
  onSyncComplete?: (data: SyncStatusData) => void;
  /** Callback when folder counts update */
  onFolderCountsUpdate?: (data: FolderCountData) => void;
  /** Callback when sync starts */
  onSyncStarted?: (data: SyncStatusData) => void;
  /** Callback when sync errors */
  onSyncError?: (data: SyncStatusData) => void;
}

interface StaffEmailWebSocketResult {
  /** Whether the WebSocket is currently connected */
  isConnected: boolean;
  /** Whether a sync is currently in progress */
  isSyncing: boolean;
  /** Manually reconnect the WebSocket */
  reconnect: () => void;
  /** Send a message to the WebSocket (for future client commands) */
  send: (data: unknown) => void;
}

/**
 * Custom hook to manage WebSocket connection for real-time staff email updates.
 * Automatically invalidates React Query cache when email changes occur.
 *
 * @example
 * ```tsx
 * const { isConnected, isSyncing } = useStaffEmailWebSocket({
 *   showToasts: true,
 *   onNewEmail: (data) => {
 *     console.log('New email:', data.subject);
 *   },
 *   onSyncComplete: () => {
 *     console.log('Email sync finished');
 *   },
 * });
 * ```
 */
export const useStaffEmailWebSocket = (
  options: UseStaffEmailWebSocketOptions = {}
): StaffEmailWebSocketResult => {
  const {
    showToasts = true,
    onNewEmail,
    onSyncComplete,
    onFolderCountsUpdate,
    onSyncStarted,
    onSyncError,
  } = options;

  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
  const reconnectAttempts = useRef(0);
  const maxReconnectAttempts = 5;
  const [isConnected, setIsConnected] = useState(false);
  const [isSyncing, setIsSyncing] = useState(false);
  const { data: user } = useCurrentUser();
  const queryClient = useQueryClient();

  const handleMessage = useCallback(
    (event: MessageEvent) => {
      try {
        const message: StaffEmailWebSocketMessage = JSON.parse(event.data);
        console.log('Staff Email WebSocket message received:', message);

        switch (message.type) {
          case 'new_email': {
            const data = message.data as NewEmailData;
            // Invalidate email list and folders
            queryClient.invalidateQueries({ queryKey: ['staff-emails'] });
            queryClient.invalidateQueries({ queryKey: ['staff-email-folders'] });
            if (showToasts) {
              toast.success(
                `New email from ${data.from_name || data.from_address || 'Unknown'}`,
                {
                  duration: 5000,
                  position: 'top-right',
                  icon: '📧',
                }
              );
            }
            onNewEmail?.(data);
            break;
          }

          case 'email_read':
          case 'email_unread':
          case 'email_moved':
          case 'email_deleted': {
            const data = message.data as NewEmailData;
            // Invalidate email list and specific email
            queryClient.invalidateQueries({ queryKey: ['staff-emails'] });
            queryClient.invalidateQueries({ queryKey: ['staff-email-folders'] });
            if (data.id) {
              queryClient.invalidateQueries({
                queryKey: ['staff-email', data.id],
              });
            }
            break;
          }

          case 'folder_counts': {
            const data = message.data as FolderCountData;
            // Update folder counts without full refetch
            queryClient.invalidateQueries({ queryKey: ['staff-email-folders'] });
            onFolderCountsUpdate?.(data);
            break;
          }

          case 'sync_started': {
            const data = message.data as SyncStatusData;
            setIsSyncing(true);
            if (showToasts) {
              toast.loading('Syncing emails...', {
                id: 'email-sync',
                position: 'bottom-right',
              });
            }
            onSyncStarted?.(data);
            break;
          }

          case 'sync_completed': {
            const data = message.data as SyncStatusData;
            setIsSyncing(false);
            // Invalidate all email-related queries
            queryClient.invalidateQueries({ queryKey: ['staff-emails'] });
            queryClient.invalidateQueries({ queryKey: ['staff-email-folders'] });
            queryClient.invalidateQueries({ queryKey: ['staff-email-addresses'] });
            if (showToasts) {
              const newCount = data.details?.new_count || data.new_count || 0;
              toast.success(
                newCount > 0 ? `Synced ${newCount} new email${newCount === 1 ? '' : 's'}` : 'Emails synced',
                {
                  id: 'email-sync',
                  duration: 3000,
                  position: 'bottom-right',
                }
              );
            }
            onSyncComplete?.(data);
            break;
          }

          case 'sync_error': {
            const data = message.data as SyncStatusData;
            setIsSyncing(false);
            if (showToasts) {
              const errorMsg = data.details?.message || data.message || 'Sync failed';
              toast.error(`Email sync error: ${errorMsg}`, {
                id: 'email-sync',
                duration: 5000,
                position: 'bottom-right',
              });
            }
            onSyncError?.(data);
            break;
          }

          default:
            console.log('Unknown staff email WebSocket message type:', message.type);
        }
      } catch (error) {
        console.error('Error parsing staff email WebSocket message:', error);
      }
    },
    [
      queryClient,
      showToasts,
      onNewEmail,
      onSyncComplete,
      onFolderCountsUpdate,
      onSyncStarted,
      onSyncError,
    ]
  );

  const connect = useCallback(() => {
    if (!user || !user.id) {
      return;
    }

    // Determine WebSocket URL using same logic as API config
    const baseDomain = getBaseDomain();
    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';

    // For localhost or lvh.me, use port 8000. In production, no port (Traefik handles it)
    const isDev = baseDomain === 'localhost' || baseDomain === 'lvh.me';
    const port = isDev ? ':8000' : '';

    const token = getCookie('access_token');
    const wsUrl = `${protocol}//api.${baseDomain}${port}/ws/staff-email/?token=${token}`;

    console.log('Connecting to staff email WebSocket:', wsUrl);

    try {
      wsRef.current = new WebSocket(wsUrl);

      wsRef.current.onopen = () => {
        console.log('Staff Email WebSocket connected');
        setIsConnected(true);
        reconnectAttempts.current = 0; // Reset reconnect attempts on successful connection
      };

      wsRef.current.onmessage = handleMessage;

      wsRef.current.onclose = (event) => {
        console.log('Staff Email WebSocket disconnected:', event.code, event.reason);
        setIsConnected(false);
        setIsSyncing(false);

        // Attempt to reconnect with exponential backoff
        if (user && user.id && reconnectAttempts.current < maxReconnectAttempts) {
          const delay = Math.min(1000 * Math.pow(2, reconnectAttempts.current), 30000);
          reconnectAttempts.current += 1;
          console.log(
            `Attempting to reconnect staff email WebSocket in ${delay}ms (attempt ${reconnectAttempts.current})`
          );

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

      wsRef.current.onerror = (error) => {
        console.error('Staff Email WebSocket error:', error);
      };
    } catch (error) {
      console.error('Failed to create staff email WebSocket:', error);
    }
  }, [user, handleMessage]);

  const reconnect = useCallback(() => {
    // Clear any existing reconnect timeout
    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current);
    }
    // Close existing connection
    if (wsRef.current) {
      wsRef.current.close();
      wsRef.current = null;
    }
    // Reset reconnect attempts
    reconnectAttempts.current = 0;
    // Connect
    connect();
  }, [connect]);

  const send = useCallback((data: unknown) => {
    if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(data));
    } else {
      console.warn('Cannot send message: WebSocket is not connected');
    }
  }, []);

  useEffect(() => {
    connect();

    return () => {
      // Clear reconnect timeout
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
      }
      // Close WebSocket
      if (wsRef.current) {
        wsRef.current.close();
        wsRef.current = null;
      }
    };
  }, [connect]);

  return {
    isConnected,
    isSyncing,
    reconnect,
    send,
  };
};

export default useStaffEmailWebSocket;