Codewalkers/apps/web/src/hooks/useSubscriptionWithErrorHandling.ts

import { useCallback, useEffect, useRef, useState } from 'react';
import { trpc } from '@/lib/trpc';
import type { SubscriptionEvent } from '@codewalk-district/shared';

interface UseSubscriptionWithErrorHandlingOptions {
  /** Called when subscription receives data */
  onData?: (data: SubscriptionEvent) => void;
  /** Called when subscription encounters an error */
  onError?: (error: Error) => void;
  /** Called when subscription starts */
  onStarted?: () => void;
  /** Called when subscription stops */
  onStopped?: () => void;
  /** Whether to automatically reconnect on errors (default: true) */
  autoReconnect?: boolean;
  /** Delay before attempting reconnection in ms (default: 1000) */
  reconnectDelay?: number;
  /** Maximum number of reconnection attempts (default: 5) */
  maxReconnectAttempts?: number;
  /** Whether the subscription is enabled (default: true) */
  enabled?: boolean;
}

interface SubscriptionState {
  isConnected: boolean;
  isConnecting: boolean;
  error: Error | null;
  reconnectAttempts: number;
  lastEventId: string | null;
}

/**
 * Hook for managing tRPC subscriptions with error handling, reconnection, and cleanup.
 *
 * Provides automatic reconnection on connection failures, tracks connection state,
 * and ensures proper cleanup on unmount.
 */
export function useSubscriptionWithErrorHandling(
  subscription: () => ReturnType<typeof trpc.onEvent.useSubscription>,
  options: UseSubscriptionWithErrorHandlingOptions = {}
) {
  const {
    autoReconnect = true,
    reconnectDelay = 1000,
    maxReconnectAttempts = 5,
    enabled = true,
  } = options;

  const [state, setState] = useState<SubscriptionState>({
    isConnected: false,
    isConnecting: false,
    error: null,
    reconnectAttempts: 0,
    lastEventId: null,
  });

  // Store callbacks in refs so they never appear in effect deps.
  // Callers pass inline arrows that change identity every render —
  // putting them in deps causes setState → re-render → new callback → effect re-fire → infinite loop.
  const callbacksRef = useRef(options);
  callbacksRef.current = options;

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

  // Clear reconnect timeout on unmount
  useEffect(() => {
    mountedRef.current = true;
    return () => {
      mountedRef.current = false;
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
      }
      callbacksRef.current.onStopped?.();
    };
  }, []);

  const scheduleReconnect = useCallback(() => {
    if (!autoReconnect || reconnectAttemptsRef.current >= maxReconnectAttempts || !mountedRef.current) {
      return;
    }

    reconnectTimeoutRef.current = setTimeout(() => {
      if (mountedRef.current) {
        reconnectAttemptsRef.current += 1;
        setState(prev => ({
          ...prev,
          isConnecting: true,
          reconnectAttempts: reconnectAttemptsRef.current,
        }));
      }
    }, reconnectDelay);
  }, [autoReconnect, maxReconnectAttempts, reconnectDelay]);

  const subscriptionResult = subscription();

  // Handle subscription state changes.
  // Only depends on primitive/stable values — never on caller callbacks.
  useEffect(() => {
    if (!enabled) {
      setState(prev => {
        if (!prev.isConnected && !prev.isConnecting && prev.error === null) return prev;
        return { ...prev, isConnected: false, isConnecting: false, error: null };
      });
      return;
    }

    if (subscriptionResult.status === 'pending') {
      setState(prev => {
        if (prev.isConnecting && prev.error === null) return prev;
        return { ...prev, isConnecting: true, error: null };
      });
      callbacksRef.current.onStarted?.();
    } else if (subscriptionResult.status === 'error') {
      const error = subscriptionResult.error instanceof Error
        ? subscriptionResult.error
        : new Error('Subscription error');

      setState(prev => ({
        ...prev,
        isConnected: false,
        isConnecting: false,
        error,
      }));

      callbacksRef.current.onError?.(error);
      scheduleReconnect();
    } else if (subscriptionResult.status === 'success') {
      reconnectAttemptsRef.current = 0;
      setState(prev => {
        if (prev.isConnected && !prev.isConnecting && prev.error === null && prev.reconnectAttempts === 0) return prev;
        return { ...prev, isConnected: true, isConnecting: false, error: null, reconnectAttempts: 0 };
      });

      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
        reconnectTimeoutRef.current = null;
      }
    }
  }, [enabled, subscriptionResult.status, subscriptionResult.error, scheduleReconnect]);

  // Handle incoming data
  useEffect(() => {
    if (subscriptionResult.data) {
      setState(prev => ({ ...prev, lastEventId: subscriptionResult.data.id }));
      callbacksRef.current.onData?.(subscriptionResult.data);
    }
  }, [subscriptionResult.data]);

  return {
    ...state,
    /** Manually trigger a reconnection attempt */
    reconnect: () => {
      if (mountedRef.current) {
        reconnectAttemptsRef.current = 0;
        setState(prev => ({
          ...prev,
          isConnecting: true,
          error: null,
          reconnectAttempts: 0,
        }));
      }
    },
    /** Reset error state and reconnection attempts */
    reset: () => {
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
        reconnectTimeoutRef.current = null;
      }
      reconnectAttemptsRef.current = 0;
      setState(prev => ({
        ...prev,
        error: null,
        reconnectAttempts: 0,
        isConnecting: false,
      }));
    },
  };
}