System Design #12: Design a Multi-Step Form Wizard

type WizardState = 
  | { status: 'filling'; currentStep: string; direction: 'forward' | 'backward' }
  | { status: 'validating'; currentStep: string }
  | { status: 'submitting' }
  | { status: 'success'; result: SubmissionResult }
  | { status: 'error'; error: string; lastStep: string };

type WizardAction =
  | { type: 'NEXT' }
  | { type: 'PREVIOUS' }
  | { type: 'JUMP_TO'; step: string }
  | { type: 'VALIDATION_SUCCESS' }
  | { type: 'VALIDATION_FAILURE'; errors: Record<string, string> }
  | { type: 'SUBMIT' }
  | { type: 'SUBMIT_SUCCESS'; result: SubmissionResult }
  | { type: 'SUBMIT_FAILURE'; error: string }
  | { type: 'SAVE_DRAFT' };

// Steps aren't always linear — they form a directed graph
interface StepDefinition {
  id: string;
  title: string;
  component: React.LazyExoticComponent<any>;
  schema: z.ZodSchema;          // Validation schema for this step
  next: (data: FormData) => string | null; // Dynamic next step
  canSkip?: boolean;
  isVisible?: (data: FormData) => boolean; // Conditional visibility
}

const steps: StepDefinition[] = [
  {
    id: 'personal-info',
    title: 'Personal Info',
    component: lazy(() => import('./steps/PersonalInfo')),
    schema: personalInfoSchema,
    next: () => 'address',
  },
  {
    id: 'address',
    title: 'Address',
    component: lazy(() => import('./steps/Address')),
    schema: addressSchema,
    next: (data) => data.needsBilling ? 'billing' : 'review', // Conditional!
  },
  {
    id: 'billing',
    title: 'Billing',
    component: lazy(() => import('./steps/Billing')),
    schema: billingSchema,
    isVisible: (data) => data.needsBilling === true, // Only show if needed
    next: () => 'review',
  },
  {
    id: 'review',
    title: 'Review',
    component: lazy(() => import('./steps/Review')),
    schema: z.object({}), // No validation on review step
    next: () => null,     // Final step
  },
];

interface WizardFormState {
  // Step data — keyed by step ID
  data: Record<string, Record<string, any>>;
  
  // Meta
  currentStepId: string;
  visitedSteps: Set<string>;      // Track which steps user has been to
  completedSteps: Set<string>;    // Steps that passed validation
  errors: Record<string, Record<string, string>>; // stepId → field → error
  
  // Persistence
  draftId: string | null;         // Server-side draft ID
  lastSavedAt: string | null;
  isDirty: boolean;               // Has unsaved changes
  
  // Files (handled separately from JSON data)
  files: Record<string, File[]>; // stepId → uploaded files
}

// Persist to sessionStorage for crash recovery
function usePersistentWizard(wizardId: string) {
  const STORAGE_KEY = `wizard-${wizardId}`;
  
  const [state, dispatch] = useReducer(wizardReducer, null, () => {
    const saved = sessionStorage.getItem(STORAGE_KEY);
    return saved ? JSON.parse(saved) : createInitialState();
  });
  
  // Auto-save on every state change
  useEffect(() => {
    sessionStorage.setItem(STORAGE_KEY, JSON.stringify({
      data: state.data,
      currentStepId: state.currentStepId,
      visitedSteps: [...state.visitedSteps],
      completedSteps: [...state.completedSteps],
    }));
  }, [state]);
  
  // Clear on successful submission
  const clearDraft = () => sessionStorage.removeItem(STORAGE_KEY);
  
  return { state, dispatch, clearDraft };
}

import { z } from 'zod';

// Per-step schemas
const personalInfoSchema = z.object({
  firstName: z.string().min(1, 'Required'),
  lastName: z.string().min(1, 'Required'),
  email: z.string().email('Invalid email'),
  phone: z.string().regex(/^\+?[1-9]\d{7,14}$/, 'Invalid phone').optional(),
});

const addressSchema = z.object({
  street: z.string().min(1, 'Required'),
  city: z.string().min(1, 'Required'),
  state: z.string().min(1, 'Required'),
  zip: z.string().min(5, 'Invalid zip code'),
  country: z.string().min(1, 'Required'),
  needsBilling: z.boolean(),
});

// Cross-step validation (runs before final submission)
const fullFormSchema = z.object({
  personalInfo: personalInfoSchema,
  address: addressSchema,
  billing: billingSchema.optional(),
}).refine(
  (data) => !(data.address.needsBilling && !data.billing),
  { message: 'Billing information required', path: ['billing'] }
);

// Async validation (debounced)
async function validateEmailUnique(email: string): Promise<string | null> {
  const { exists } = await api.get(`/check-email?email=${email}`);
  return exists ? 'Email already registered' : null;
}

function useWizard(config: WizardConfig) {
  return {
    // Navigation
    currentStep: StepDefinition;
    currentStepIndex: number;
    totalSteps: number;
    visibleSteps: StepDefinition[];  // After conditional filtering
    goNext: () => Promise<boolean>;  // Returns false if validation fails
    goPrevious: () => void;
    jumpTo: (stepId: string) => void;
    canGoNext: boolean;
    canGoPrevious: boolean;
    isFirstStep: boolean;
    isLastStep: boolean;
    
    // Data
    stepData: Record<string, any>;   // Current step's data
    setStepData: (field: string, value: any) => void;
    allData: Record<string, Record<string, any>>;
    
    // Validation
    errors: Record<string, string>;  // Current step's errors
    isValidating: boolean;
    
    // Submission
    submit: () => Promise<void>;
    isSubmitting: boolean;
    
    // Persistence
    saveDraft: () => Promise<void>;
    isDirty: boolean;
    lastSavedAt: string | null;
    
    // Progress
    progress: number;               // 0-100
    visitedSteps: Set<string>;
    completedSteps: Set<string>;
  };
}

interface StepComponentProps {
  data: Record<string, any>;
  onChange: (field: string, value: any) => void;
  errors: Record<string, string>;
  allFormData: Record<string, Record<string, any>>; // Access other steps' data
}

// Each step is a standalone form section
function PersonalInfoStep({ data, onChange, errors }: StepComponentProps) {
  return (
    <div className="space-y-4">
      <div>
        <Label htmlFor="firstName">First Name *</Label>
        <Input
          id="firstName"
          value={data.firstName || ''}
          onChange={(e) => onChange('firstName', e.target.value)}
          aria-invalid={!!errors.firstName}
          aria-describedby={errors.firstName ? 'firstName-error' : undefined}
        />
        {errors.firstName && (
          <p id="firstName-error" className="text-sm text-destructive mt-1">
            {errors.firstName}
          </p>
        )}
      </div>
      {/* ... more fields */}
    </div>
  );
}

function StepContainer({ step, direction }: { step: StepDefinition; direction: 'forward' | 'backward' }) {
  const headingRef = useRef<HTMLHeadingElement>(null);
  
  useEffect(() => {
    // Focus the step heading on transition
    headingRef.current?.focus();
    
    // Announce to screen readers
    announce(`Step ${stepIndex + 1} of ${totalSteps}: ${step.title}`);
  }, [step.id]);
  
  return (
    <motion.div
      key={step.id}
      initial={{ x: direction === 'forward' ? 100 : -100, opacity: 0 }}
      animate={{ x: 0, opacity: 1 }}
      exit={{ x: direction === 'forward' ? -100 : 100, opacity: 0 }}
      transition={{ duration: 0.2 }}
    >
      <h2 ref={headingRef} tabIndex={-1} className="text-lg font-semibold outline-none">
        {step.title}
      </h2>
      <Suspense fallback={<StepSkeleton />}>
        <step.component {...stepProps} />
      </Suspense>
    </motion.div>
  );
}

function WizardProgress({ steps, currentIndex, completedSteps, visitedSteps }: WizardProgressProps) {
  return (
    <nav aria-label="Form progress">
      <ol className="flex items-center gap-2">
        {steps.map((step, i) => {
          const isCompleted = completedSteps.has(step.id);
          const isCurrent = i === currentIndex;
          const isClickable = visitedSteps.has(step.id) || isCompleted;
          
          return (
            <li key={step.id} className="flex items-center">
              <button
                onClick={() => isClickable && jumpTo(step.id)}
                disabled={!isClickable}
                aria-current={isCurrent ? 'step' : undefined}
                className={cn(
                  "w-8 h-8 rounded-full flex items-center justify-center text-xs font-medium",
                  isCompleted && "bg-primary text-primary-foreground",
                  isCurrent && "border-2 border-primary text-primary",
                  !isCompleted && !isCurrent && "bg-muted text-muted-foreground"
                )}
              >
                {isCompleted ? <Check className="h-4 w-4" /> : i + 1}
              </button>
              {i < steps.length - 1 && (
                <div className={cn("w-12 h-0.5 mx-1", isCompleted ? "bg-primary" : "bg-muted")} />
              )}
            </li>
          );
        })}
      </ol>
    </nav>
  );
}

function useAutoSaveDraft(wizardId: string, formData: Record<string, any>, userId: string | null) {
  const debouncedSave = useMemo(
    () => debounce(async (data: Record<string, any>) => {
      if (!userId) return; // Only save for authenticated users
      
      await api.put(`/drafts/${wizardId}`, {
        data,
        updated_at: new Date().toISOString(),
      });
    }, 3000), // Save 3s after last change
    [wizardId, userId]
  );
  
  useEffect(() => {
    debouncedSave(formData);
    return () => debouncedSave.cancel();
  }, [formData, debouncedSave]);
}

function useStepAnalytics(wizardId: string, currentStepId: string) {
  const stepStartTime = useRef(Date.now());
  
  // Track time spent on each step
  useEffect(() => {
    stepStartTime.current = Date.now();
    
    return () => {
      const timeSpent = Date.now() - stepStartTime.current;
      analytics.track('wizard_step_time', {
        wizardId,
        stepId: currentStepId,
        timeSpentMs: timeSpent,
      });
    };
  }, [currentStepId]);
  
  // Track drop-off (user leaves without completing)
  useEffect(() => {
    const handleBeforeUnload = () => {
      analytics.track('wizard_abandon', {
        wizardId,
        lastStepId: currentStepId,
        // This tells you WHERE users give up
      });
    };
    
    window.addEventListener('beforeunload', handleBeforeUnload);
    return () => window.removeEventListener('beforeunload', handleBeforeUnload);
  }, [currentStepId]);
}

function useWizardHistory(currentStepId: string, goPrevious: () => void) {
  useEffect(() => {
    // Push a history entry for each step
    window.history.pushState({ step: currentStepId }, '');
    
    const handlePopState = (e: PopStateEvent) => {
      if (e.state?.step) {
        goPrevious();
      }
    };
    
    window.addEventListener('popstate', handlePopState);
    return () => window.removeEventListener('popstate', handlePopState);
  }, [currentStepId]);
}