@matthew.ngo/reform/src/features/wizard/useFormWizard.ts

A flexible and powerful React form management library with advanced validation, state observation, and multi-group support

import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
import { ReformReturn } from '../../types';
import {
  FormStep,
  FormWizardConfig,
  FormWizardReturn,
  FormWizardState,
} from './types';
import { useMemoizedCallback } from '../../common/useMemoizedCallback';

/**
 * Hook for creating and managing multi-step form wizards
 *
 * @template T - The type of form data
 * @param reform - The Reform hook return value
 * @param config - Configuration for the form wizard
 * @returns Form wizard state and methods
 */
export const useFormWizard = <T extends Record<string, any>>(
  reform: ReformReturn<T>,
  config: FormWizardConfig<T>
): FormWizardReturn<T> => {
  // Memoize config to prevent unnecessary re-renders
  const memoizedConfig = useMemo(
    () => ({
      steps: config.steps,
      initialStepId: config.initialStepId,
      validateOnNext: config.validateOnNext ?? true,
      allowSkipSteps: config.allowSkipSteps ?? false,
      markCompletedOnLeave: config.markCompletedOnLeave ?? true,
      persistState: config.persistState ?? false,
      persistStateKey: config.persistStateKey ?? 'reform-wizard-state',
      onStepCompleted: config.onStepCompleted,
      onWizardCompleted: config.onWizardCompleted,
    }),
    [
      // Use JSON.stringify for complex objects to compare by value
      JSON.stringify(config.steps),
      config.initialStepId,
      config.validateOnNext,
      config.allowSkipSteps,
      config.markCompletedOnLeave,
      config.persistState,
      config.persistStateKey,
      // Functions are compared by reference
      config.onStepCompleted,
      config.onWizardCompleted,
    ]
  );

  const {
    steps,
    initialStepId,
    validateOnNext,
    allowSkipSteps,
    markCompletedOnLeave,
    persistState,
    persistStateKey,
    onStepCompleted,
    onWizardCompleted,
  } = memoizedConfig;

  // Store the reform reference to avoid unnecessary re-renders
  const reformRef = useRef(reform);
  useEffect(() => {
    reformRef.current = reform;
  }, [reform]);

  // Initialize state from storage if enabled
  const initialState = useMemo((): FormWizardState => {
    if (persistState) {
      try {
        const savedState = localStorage.getItem(persistStateKey);
        if (savedState) {
          const parsedState = JSON.parse(savedState) as FormWizardState;
          return parsedState;
        }
      } catch (error) {
        console.error('Error loading wizard state from storage:', error);
      }
    }

    return {
      currentStepId: initialStepId || steps[0]?.id || '',
      completedSteps: [],
      isValidating: false,
      isTransitioning: false,
      isCompleted: false,
    };
  }, [persistState, persistStateKey, initialStepId, steps]);

  // State
  const [currentStepId, setCurrentStepId] = useState<string>(
    initialState.currentStepId
  );
  const [completedSteps, setCompletedSteps] = useState<string[]>(
    initialState.completedSteps
  );
  const [isValidating, setIsValidating] = useState<boolean>(false);
  const [isTransitioning, setIsTransitioning] = useState<boolean>(false);

  // Derived state
  const isCompleted = useMemo(() => {
    return steps.every(step => completedSteps.includes(step.id));
  }, [steps, completedSteps]);

  // Persist state to storage when it changes
  useEffect(() => {
    if (persistState) {
      const stateToSave: FormWizardState = {
        currentStepId,
        completedSteps,
        isValidating,
        isTransitioning,
        isCompleted,
      };
      localStorage.setItem(persistStateKey, JSON.stringify(stateToSave));
    }
  }, [
    currentStepId,
    completedSteps,
    isValidating,
    isTransitioning,
    isCompleted,
    persistState,
    persistStateKey,
  ]);

  // Get current step
  const currentStep = useMemo(() => {
    return steps.find(step => step.id === currentStepId) || steps[0];
  }, [steps, currentStepId]);

  // Get groups for a specific step - memoized to prevent unnecessary re-renders
  const getStepGroups = useMemoizedCallback((step: FormStep<T>) => {
    const allGroups = reformRef.current.getGroups();
    return step.groupIndices.map(index => allGroups[index]).filter(Boolean);
  }, []);

  // Validate a step - memoized to prevent unnecessary re-renders
  const validateStep = useMemoizedCallback(
    async (stepId: string): Promise<boolean> => {
      const step = steps.find(s => s.id === stepId);
      if (!step) return false;

      setIsValidating(true);

      try {
        const stepGroups = getStepGroups(step);

        // If step has custom validation, use it
        if (step.validate) {
          return step.validate(stepGroups);
        }

        // Otherwise validate all groups in the step
        const groupIndices = step.groupIndices;
        // FIXME
        // const validationResults = await Promise.all(
        //   groupIndices.map(index => reformRef.current.validateGroup(index))
        // );

        // return validationResults.every(result => result);
        return await Promise.resolve(true);
      } finally {
        setIsValidating(false);
      }
    },
    [steps, getStepGroups]
  );

  // Check if a step is completed - memoized to prevent unnecessary re-renders
  const isStepCompleted = useMemoizedCallback(
    (stepId: string): boolean => {
      // If already marked as completed, return true
      if (completedSteps.includes(stepId)) return true;

      // Otherwise check custom completion logic
      const step = steps.find(s => s.id === stepId);
      if (!step) return false;

      if (step.isCompleted) {
        const stepGroups = getStepGroups(step);
        return step.isCompleted(stepGroups);
      }

      return false;
    },
    [completedSteps, steps, getStepGroups]
  );

  // Check if a step is enabled - memoized to prevent unnecessary re-renders
  const isStepEnabled = useMemoizedCallback(
    (stepId: string): boolean => {
      // If allowing skipping steps, all steps are enabled
      if (allowSkipSteps) return true;

      const step = steps.find(s => s.id === stepId);
      if (!step) return false;

      // If step has custom enabled logic, use it
      if (step.isEnabled) {
        const allGroups = reformRef.current.getGroups();
        return step.isEnabled(allGroups, completedSteps);
      }

      // Otherwise, a step is enabled if it's the first step or all previous steps are completed
      const stepIndex = steps.findIndex(s => s.id === stepId);
      if (stepIndex === 0) return true;

      const previousSteps = steps.slice(0, stepIndex);
      return previousSteps.every(s => completedSteps.includes(s.id));
    },
    [steps, allowSkipSteps, completedSteps]
  );

  // Mark a step as completed - memoized to prevent unnecessary re-renders
  const markStepCompleted = useMemoizedCallback(
    (stepId: string) => {
      if (!completedSteps.includes(stepId)) {
        const newCompletedSteps = [...completedSteps, stepId];
        setCompletedSteps(newCompletedSteps);

        // Call onStepCompleted callback
        if (onStepCompleted) {
          const step = steps.find(s => s.id === stepId);
          if (step) {
            const allGroups = reformRef.current.getGroups();
            onStepCompleted(stepId, allGroups);
          }
        }
        // Check if all steps are now completed
        const allCompleted = steps.every(step =>
          newCompletedSteps.includes(step.id)
        );

        if (allCompleted && onWizardCompleted) {
          const allGroups = reformRef.current.getGroups();
          onWizardCompleted(allGroups);
        }
      }
    },
    [steps, completedSteps, onStepCompleted, onWizardCompleted]
  );

  // Mark a step as not completed - memoized to prevent unnecessary re-renders
  const markStepNotCompleted = useMemoizedCallback(
    (stepId: string) => {
      if (completedSteps.includes(stepId)) {
        setCompletedSteps(completedSteps.filter(id => id !== stepId));
      }
    },
    [completedSteps]
  );

  // Navigation methods - memoized to prevent unnecessary re-renders
  const goToStep = useMemoizedCallback(
    async (stepId: string): Promise<boolean> => {
      const targetStep = steps.find(s => s.id === stepId);
      if (!targetStep) return false;

      // Check if step is enabled
      if (!isStepEnabled(stepId)) return false;

      // Set transitioning state
      setIsTransitioning(true);

      try {
        // Run onLeave for current step if it exists
        if (currentStep && currentStep.onLeave) {
          const direction =
            steps.findIndex(s => s.id === stepId) >
            steps.findIndex(s => s.id === currentStepId)
              ? 'next'
              : 'prev';

          const allGroups = reformRef.current.getGroups();
          const canLeave = await Promise.resolve(
            currentStep.onLeave(allGroups, direction)
          );

          if (canLeave === false) {
            return false;
          }
        }

        // Mark current step as completed if configured to do so
        if (markCompletedOnLeave && currentStepId) {
          const isValid = await validateStep(currentStepId);
          if (isValid) {
            markStepCompleted(currentStepId);
          }
        }

        // Run onEnter for target step if it exists
        if (targetStep.onEnter) {
          const allGroups = reformRef.current.getGroups();
          await Promise.resolve(targetStep.onEnter(allGroups));
        }

        // Update current step
        setCurrentStepId(stepId);
        return true;
      } finally {
        setIsTransitioning(false);
      }
    },
    [
      steps,
      currentStep,
      currentStepId,
      isStepEnabled,
      markCompletedOnLeave,
      validateStep,
      markStepCompleted,
    ]
  );

  // Navigate to next step - memoized to prevent unnecessary re-renders
  const next = useMemoizedCallback(async (): Promise<boolean> => {
    // Find current step index
    const currentIndex = steps.findIndex(s => s.id === currentStepId);
    if (currentIndex === -1 || currentIndex >= steps.length - 1) return false;

    // Validate current step if required
    if (validateOnNext) {
      const isValid = await validateStep(currentStepId);
      if (!isValid) return false;

      // Mark current step as completed if valid
      markStepCompleted(currentStepId);
    }

    // Go to next step
    const nextStep = steps[currentIndex + 1];
    return goToStep(nextStep.id);
  }, [
    steps,
    currentStepId,
    validateOnNext,
    validateStep,
    markStepCompleted,
    goToStep,
  ]);

  // Navigate to previous step - memoized to prevent unnecessary re-renders
  const prev = useMemoizedCallback(async (): Promise<boolean> => {
    // Find current step index
    const currentIndex = steps.findIndex(s => s.id === currentStepId);
    if (currentIndex <= 0) return false;

    // Go to previous step
    const prevStep = steps[currentIndex - 1];
    return goToStep(prevStep.id);
  }, [steps, currentStepId, goToStep]);

  // Reset wizard - memoized to prevent unnecessary re-renders
  const resetWizard = useMemoizedCallback(
    (options?: { resetFormData?: boolean; clearCompletedSteps?: boolean }) => {
      const { resetFormData = false, clearCompletedSteps = true } =
        options || {};

      // Reset to initial step
      setCurrentStepId(initialStepId || steps[0]?.id || '');

      // Clear completed steps if requested
      if (clearCompletedSteps) {
        setCompletedSteps([]);
      }

      // Reset form data if requested
      if (resetFormData) {
        reformRef.current.formMethods.reset();
      }
    },
    [initialStepId, steps]
  );

  // Get all steps with status information - memoized to prevent unnecessary re-renders
  const getSteps = useMemoizedCallback(() => {
    return steps.map(step => {
      const {
        isEnabled: checkEnabled,
        isCompleted: checkCompleted,
        ...restStep
      } = step;

      return {
        ...restStep,
        isActive: step.id === currentStepId,
        isCompleted: isStepCompleted(step.id),
        isEnabled: isStepEnabled(step.id),
        checkEnabled,
        checkCompleted,
      };
    });
  }, [steps, currentStepId, isStepCompleted, isStepEnabled]);

  // Get a specific step with status information - memoized to prevent unnecessary re-renders
  const getStep = useMemoizedCallback(
    (stepId: string) => {
      const step = steps.find(s => s.id === stepId);
      if (!step) return undefined;

      const {
        isEnabled: checkEnabled,
        isCompleted: checkCompleted,
        ...restStep
      } = step;

      return {
        ...restStep,
        isActive: step.id === currentStepId,
        isCompleted: isStepCompleted(step.id),
        isEnabled: isStepEnabled(step.id),
        checkEnabled,
        checkCompleted,
      };
    },
    [steps, currentStepId, isStepCompleted, isStepEnabled]
  );

  // Get current step with status information - memoized to prevent unnecessary re-renders
  const getCurrentStep = useMemoizedCallback(() => {
    const step = steps.find(s => s.id === currentStepId) || steps[0];

    const {
      isEnabled: checkEnabled,
      isCompleted: checkCompleted,
      ...restStep
    } = step;

    return {
      ...restStep,
      isActive: true,
      isCompleted: isStepCompleted(step.id),
      isEnabled: isStepEnabled(step.id),
      checkEnabled,
      checkCompleted,
    };
  }, [steps, currentStepId, isStepCompleted, isStepEnabled]);

  // Add utility methods for step navigation - memoized to prevent unnecessary re-renders
  const canGoNext = useMemoizedCallback((): boolean => {
    const currentIndex = steps.findIndex(s => s.id === currentStepId);
    return currentIndex < steps.length - 1;
  }, [steps, currentStepId]);

  const canGoPrev = useMemoizedCallback((): boolean => {
    const currentIndex = steps.findIndex(s => s.id === currentStepId);
    return currentIndex > 0;
  }, [steps, currentStepId]);

  // Get step index - memoized to prevent unnecessary re-renders
  const getStepIndex = useMemoizedCallback(
    (stepId: string): number => {
      return steps.findIndex(s => s.id === stepId);
    },
    [steps]
  );

  // Get current step index - memoized to prevent unnecessary re-renders
  const getCurrentStepIndex = useMemoizedCallback((): number => {
    return steps.findIndex(s => s.id === currentStepId);
  }, [steps, currentStepId]);

  // Jump to a specific step by index - memoized to prevent unnecessary re-renders
  const goToStepByIndex = useMemoizedCallback(
    async (index: number): Promise<boolean> => {
      if (index < 0 || index >= steps.length) return false;
      return goToStep(steps[index].id);
    },
    [steps, goToStep]
  );

  // Check if a step is the current step - memoized to prevent unnecessary re-renders
  const isCurrentStep = useMemoizedCallback(
    (stepId: string): boolean => {
      return stepId === currentStepId;
    },
    [currentStepId]
  );

  // Memoize the return object to prevent unnecessary re-renders
  return useMemo(
    () => ({
      // Step data
      steps,
      currentStep,
      currentStepId,
      completedSteps,
      isCompleted,
      isValidating,
      isTransitioning,

      // Step status methods
      isStepCompleted,
      isStepEnabled,
      isCurrentStep,

      // Step management
      markStepCompleted,
      markStepNotCompleted,
      validateStep,

      // Navigation methods
      goToStep,
      next,
      prev,
      canGoNext: canGoNext(),
      canGoPrev: canGoPrev(),
      resetWizard,

      // Step utilities
      getSteps,
      getStep,
      getCurrentStep,
      getStepGroups,
      getStepIndex,
      getCurrentStepIndex,
      goToStepByIndex,

      // Config
      config: memoizedConfig,
    }),
    [
      steps,
      currentStep,
      currentStepId,
      completedSteps,
      isCompleted,
      isValidating,
      isTransitioning,
      isStepCompleted,
      isStepEnabled,
      isCurrentStep,
      markStepCompleted,
      markStepNotCompleted,
      validateStep,
      goToStep,
      next,
      prev,
      canGoNext,
      canGoPrev,
      resetWizard,
      getSteps,
      getStep,
      getCurrentStep,
      getStepGroups,
      getStepIndex,
      getCurrentStepIndex,
      goToStepByIndex,
      memoizedConfig,
    ]
  );
};