Multi-Step Flow

"use client";

import * as React from "react";
import {
  AnimatePresence,
  motion,
  MotionConfig,
  useReducedMotion,
  type HTMLMotionProps,
} from "motion/react";

import { Button } from "@/components/ui/button";
import { useElementHeight } from "@/hooks/use-element-height";
import { cn } from "@/lib/utils";

type MotionTransition = NonNullable<HTMLMotionProps<"div">["transition"]>;
type StepDirection = -1 | 1;

const useIsomorphicLayoutEffect =
  typeof window === "undefined" ? React.useEffect : React.useLayoutEffect;

export type MultiStepItem = {
  id: string;
  content: React.ReactNode;
};

export type MultiStepProps = Omit<
  HTMLMotionProps<"div">,
  | "animate"
  | "children"
  | "defaultValue"
  | "initial"
  | "onChange"
  | "transition"
> & {
  steps: MultiStepItem[];
  step?: number;
  defaultStep?: number;
  onStepChange?: (step: number) => void;
  /** Called when the primary action is pressed on the last step. */
  onComplete?: () => void;
  backLabel?: React.ReactNode;
  continueLabel?: React.ReactNode;
  completeLabel?: React.ReactNode;
  disableBack?: boolean;
  disableContinue?: boolean;
  footer?: React.ReactNode;
  contentClassName?: string;
  innerClassName?: string;
  actionsClassName?: string;
  transition?: MotionTransition;
};

export function MultiStep({
  steps,
  step,
  defaultStep = 0,
  onStepChange,
  onComplete,
  backLabel = "Back",
  continueLabel = "Continue",
  completeLabel = "Done",
  disableBack,
  disableContinue,
  footer,
  className,
  contentClassName,
  innerClassName,
  actionsClassName,
  transition,
  ...props
}: MultiStepProps) {
  const [uncontrolledStep, setUncontrolledStep] = React.useState(defaultStep);
  const [innerRef, height, measureHeight] =
    useElementHeight<HTMLDivElement>();
  const shouldReduceMotion = useReducedMotion();

  const selectedStep = clampStep(step ?? uncontrolledStep, steps.length);
  const isControlled = step !== undefined;
  const isFirstStep = selectedStep === 0;
  const isLastStep = selectedStep === steps.length - 1;
  const activeStep = steps[selectedStep];

  const [transitionState, setTransitionState] = React.useState<{
    step: number;
    direction: StepDirection;
  }>({ step: selectedStep, direction: 1 });
  let direction = transitionState.direction;

  // Built-in actions update this state in the same event as the step. This
  // fallback only derives direction for externally driven controlled changes.
  if (transitionState.step !== selectedStep) {
    direction = selectedStep < transitionState.step ? -1 : 1;
    setTransitionState({ step: selectedStep, direction });
  }

  // Measure the new step before paint so the content, footer projection, and
  // container height begin from the same animation frame.
  useIsomorphicLayoutEffect(() => {
    measureHeight();
  }, [activeStep?.id, measureHeight]);

  const setStep = React.useCallback(
    (nextStep: number) => {
      const resolvedStep = clampStep(nextStep, steps.length);
      const nextDirection: StepDirection =
        resolvedStep < selectedStep ? -1 : 1;

      setTransitionState({
        step: resolvedStep,
        direction: nextDirection,
      });

      if (!isControlled) {
        setUncontrolledStep(resolvedStep);
      }

      onStepChange?.(resolvedStep);
    },
    [isControlled, onStepChange, selectedStep, steps.length],
  );

  const resolvedTransition: MotionTransition = shouldReduceMotion
    ? { duration: 0 }
    : (transition ?? { type: "spring", duration: 0.3, bounce: 0 });
  const contentVariants = shouldReduceMotion
    ? reducedMotionStepVariants
    : stepVariants;

  if (!activeStep) {
    return null;
  }

  return (
    <MotionConfig reducedMotion="user" transition={resolvedTransition}>
      {/*
        `contain: layout` scopes the per-frame reflow of the height animation to this
        element, so the layout cost stays bounded no matter how heavy a step's content
        is. This keeps the transition smooth on lower-powered / mobile devices; removing
        it can cause frame drops there. It does not affect the measured `auto` height.
      */}
      <motion.div
        {...props}
        data-slot="multi-step"
        initial={false}
        animate={
          shouldReduceMotion ? { height: "auto" } : { height: height ?? "auto" }
        }
        className={cn(
          "overflow-hidden rounded-lg border bg-background contain-[layout]",
          className,
        )}
      >
        <div
          ref={innerRef}
          data-slot="multi-step-inner"
          className={cn("flex flex-col", innerClassName)}
        >
          <div
            data-slot="multi-step-viewport"
            className="relative overflow-hidden"
          >
            <AnimatePresence
              mode={shouldReduceMotion ? "sync" : "popLayout"}
              initial={false}
              custom={direction}
            >
              <motion.div
                key={activeStep.id}
                data-slot="multi-step-content"
                custom={direction}
                variants={contentVariants}
                initial="initial"
                animate="active"
                exit="exit"
                className={cn("w-full p-4", contentClassName)}
              >
                {activeStep.content}
              </motion.div>
            </AnimatePresence>
          </div>
          {footer !== undefined ? (
            footer
          ) : (
            <motion.div
              layout={!shouldReduceMotion}
              data-slot="multi-step-actions"
              className={cn(
                "flex items-center justify-between gap-3 border-t p-4",
                actionsClassName,
              )}
            >
              <Button
                type="button"
                variant="outline"
                disabled={disableBack || isFirstStep}
                onClick={() => setStep(selectedStep - 1)}
              >
                {backLabel}
              </Button>
              <Button
                type="button"
                disabled={disableContinue || (isLastStep && !onComplete)}
                onClick={() =>
                  isLastStep ? onComplete?.() : setStep(selectedStep + 1)
                }
              >
                {/* An invisible twin of the other label reserves the wider
                    width, so the Continue ↔ Done swap on the last step
                    doesn't snap the button size. */}
                <span className="grid justify-items-center">
                  <span aria-hidden="true" className="invisible col-start-1 row-start-1">
                    {isLastStep ? continueLabel : completeLabel}
                  </span>
                  <span className="col-start-1 row-start-1">
                    {isLastStep ? completeLabel : continueLabel}
                  </span>
                </span>
              </Button>
            </motion.div>
          )}
        </div>
      </motion.div>
    </MotionConfig>
  );
}

const stepVariants = {
  initial: (direction: StepDirection) => ({
    x: `${110 * direction}%`,
    opacity: 0,
  }),
  active: { x: "0%", opacity: 1 },
  exit: (direction: StepDirection) => ({
    x: `${-110 * direction}%`,
    opacity: 0,
  }),
};

const reducedMotionStepVariants = {
  initial: { opacity: 0 },
  active: { opacity: 1 },
  exit: { opacity: 0 },
};

function clampStep(step: number, stepCount: number) {
  if (stepCount <= 0) return 0;

  return Math.min(Math.max(step, 0), stepCount - 1);
}