alert/src/models/useSSE.ts

import {
  EventStreamContentType,
  fetchEventSource,
} from '@microsoft/fetch-event-source';
import { useCallback, useRef, useState } from 'react';
import { authFaiedToLoginPage } from '@/access';
export type SSEStatus =
  | 'idle'
  | 'connecting'
  | 'open'
  | 'error'
  | 'closed'
  | 'replaced';
interface SSEOptions {
  method?: 'GET' | 'POST';
  headers?: Record<string, string>;
  body?: any;
  maxRetries?: number;
  retryInterval?: number;
  openWhenHidden?: boolean;
  onMessage?: (data: string) => void;
  onError?: (error: any) => void;
  onStatusChange?: (status: SSEStatus) => void;
}

export default () => {
  const [status, setStatus] = useState<SSEStatus>('idle');
  const [retryCount, setRetryCount] = useState(0);

  // Use ref to synchronize internal logic to avoid logical judgment errors caused by status update delays.
  const internalCountRef = useRef(0);
  const abortControllerRef = useRef<AbortController | null>(null);

  // Save retry parameters for onclose use
  const retryParamsRef = useRef<{
    url: string;
    options: SSEOptions;
    maxRetries: number;
    retryInterval: number;
  } | null>(null);

  const updateStatus = useCallback((s: SSEStatus) => {
    setStatus(s);
  }, []);

  /**
   * Execute retry logic
   * @returns true means you can continue to retry, false means the maximum number of retries has been reached
   */
  const doRetry = useCallback(
    (err?: any): boolean => {
      const { maxRetries = 10, options } = retryParamsRef.current || {};

      if (!retryParamsRef.current) return false;

      internalCountRef.current += 1;
      setRetryCount(internalCountRef.current);

      // Maximum number of retries reached
      if (internalCountRef.current >= maxRetries) {
        updateStatus('error');
        options?.onError?.(
          err || new Error('SSE connection failed after max retries'),
        );
        return false;
      }
      // Continue retrying
      updateStatus('connecting');
      return true;
    },
    [updateStatus],
  );

  const connectSSE = useCallback(
    (url: string, options: SSEOptions) => {
      const {
        method = 'GET',
        headers,
        body,
        maxRetries = 3,
        retryInterval = 3000,
        openWhenHidden = true,
        onMessage,
      } = options;

      // Save retry parameters
      retryParamsRef.current = { url, options, maxRetries, retryInterval };

      // Clean up old connections
      if (abortControllerRef.current) {
        abortControllerRef.current.abort();
      }

      abortControllerRef.current = new AbortController();
      // Reset count when first connecting (if not a retry scenario)
      if (status === 'idle' || status === 'error' || status === 'closed') {
        internalCountRef.current = 0;
        setRetryCount(0);
      }
      updateStatus('connecting');

      fetchEventSource(url, {
        method,
        headers: { ...headers },
        body: body ? JSON.stringify(body) : undefined,
        signal: abortControllerRef.current.signal,
        openWhenHidden,
        async onopen(response) {
          if (
            response.ok &&
            response.headers
              .get('content-type')
              ?.includes(EventStreamContentType)
          ) {
            internalCountRef.current = 0;
            setRetryCount(0);
            updateStatus('open');
            return;
          }

          // 401 Handling: Clear token and redirect to login page
          if (response.status === 401) {
            await authFaiedToLoginPage();
            return;
          }

          throw new Error(`Connection failed with status ${response.status}`);
        },

        onmessage(msg) {
          const event = msg.event.toLowerCase();
          // Filter notification events
          if (msg.data && event === 'data') {
            onMessage?.(msg.data);
          }

          /**
           * Other clients log in and connect to SSE. The current logic only makes one connection per user.
           * The previous connection will wait for the default timeout in the background before disconnecting and will not trigger a disconnect event after testing.
           * So it is agreed to actively disconnect the connection with this event and notify the user, allowing them to decide whether to reconnect.
           */
          if (event === 'replaced') {
            disconnectSSE('replaced');
          }

          if (event === 'error') {
            console.log('backend notice see closed');
            const shouldRetry = doRetry();
            if (shouldRetry) {
              console.log(
                `Retrying connection in ${retryInterval}ms... (attempt ${internalCountRef.current})`,
              );
              // Use setTimeout to delay retries
              setTimeout(() => {
                // Check if actively disconnected
                if (abortControllerRef.current) {
                  connectSSE(url, options);
                }
              }, retryInterval);
            } else {
              updateStatus('error');
            }
          }
        },

        onclose() {
          console.log('SSE connection closed');
          // Try to reconnect
          const shouldRetry = doRetry();

          if (shouldRetry) {
            console.log(
              `Retrying connection in ${retryInterval}ms... (attempt ${internalCountRef.current})`,
            );
            // Use setTimeout to delay retries
            setTimeout(() => {
              // Check if actively disconnected
              if (abortControllerRef.current) {
                connectSSE(url, options);
              }
            }, retryInterval);
          } else {
            updateStatus('error');
          }
        },

        onerror(err) {
          console.error('SSE connection error:');

          // User actively disconnects, do not count
          if (err.name === 'AbortError') {
            updateStatus('closed');
            throw err;
          }

          const shouldRetry = doRetry(err);

          if (!shouldRetry) {
            throw err; // Maximum number of retries reached, throw exception to stop automatic reconnection
          }

          updateStatus('connecting'); // Mark as reconnecting
          return retryInterval;
        },
      });
    },
    [updateStatus, doRetry],
  );

  const disconnectSSE = useCallback(
    (status: SSEStatus = 'error') => {
      console.log('disconnect ...');

      // Clear the retry parameters to prevent continued retries in onclose
      retryParamsRef.current = null;

      if (abortControllerRef.current) {
        abortControllerRef.current.abort();
        abortControllerRef.current = null;
        updateStatus(status);
      }
    },
    [updateStatus],
  );

  return {
    status,
    retryCount,
    connectSSE,
    disconnectSSE,
  };
};