StepModal.tsx

import {
  Fragment,
  ReactNode,
  forwardRef,
  createContext,
  useContext,
  useReducer,
  useEffect,
  Children,
} from "react"
import clsx from "clsx"
import { Dialog, Transition } from "@headlessui/react"

const StepModalContext = createContext({
  currentStep: 0,
  previousStep: 0,
})

// Display a series of modal panels that can be navigated
// between with animated left / right transitions.
interface StepModalProps {
  open: boolean
  onClose: () => void
  step: number
  children?: ReactNode
}
export default function StepModal({ open, onClose, step, children }: StepModalProps) {
  const isChildrenArray = Array.isArray(children)
  const { currentStep, previousStep, set } = useSteps(isChildrenArray ? children?.length : 0)
  // We track step changes with useEffect so we can keep track of
  // the previous step and transition in the right direction.
  useEffect(() => {
    set(step)
  }, [set, step])
  return (
    <StepModalContext.Provider value={{ currentStep, previousStep }}>
      <Transition show={open} as={Fragment}>
        <Dialog
          onClose={onClose}
          className="fixed z-10 inset-0 overflow-y-auto max-w-full overflow-x-hidden [scrollbar-gutter:stable_both-edges]"
        >
          <div className="flex items-start justify-center min-h-screen">
            <Transition.Child
              as={Fragment}
              enter="ease-out duration-300"
              enterFrom="opacity-0"
              enterTo="opacity-100"
              leave="ease-in duration-200"
              leaveFrom="opacity-100"
              leaveTo="opacity-0"
            >
              <Dialog.Overlay className="fixed inset-0 bg-gray-900/30" />
            </Transition.Child>
            <Transition.Child
              as={TransitionChild}
              enter="ease-out duration-300"
              enterFrom="opacity-0 scale-95"
              enterTo="opacity-100 scale-100"
              leave="ease-in duration-200"
              leaveFrom="opacity-100 scale-100"
              leaveTo="opacity-0 scale-95"
            >
              {Children.map(children, (child, index) => (
                <ModalTransition step={index}>{child}</ModalTransition>
              ))}
            </Transition.Child>
          </div>
        </Dialog>
      </Transition>
    </StepModalContext.Provider>
  )
}

// Needed so that Transition.Child can pass a ref to a node,
// while still displaying the children in the correct way
const TransitionChild = forwardRef<HTMLDivElement>(function transitionChild(props, ref) {
  return (
    <div ref={ref} className="w-full h-full flex justify-center" {...props}>
      {props.children}
    </div>
  )
})

interface ModalTransitionProps {
  step: number
  children?: ReactNode
}
function ModalTransition({ step, children }: ModalTransitionProps) {
  const { currentStep, previousStep } = useContext(StepModalContext)
  return (
    <Transition
      as={Fragment}
      show={step === currentStep}
      enter="ease-in-out duration-300"
      enterFrom={clsx(
        "opacity-0",
        currentStep < previousStep && "-translate-x-[100vw]",
        currentStep > previousStep && "translate-x-[100vw]",
      )}
      enterTo="opacity-100 translate-x-0"
      leave="ease-in-out duration-300"
      leaveFrom="opacity-100 translate-x-0"
      leaveTo={clsx(
        "opacity-0",
        currentStep < previousStep && "translate-x-[100vw]",
        currentStep >= previousStep && "-translate-x-[100vw]",
      )}
    >
      <div className="absolute">{children}</div>
    </Transition>
  )
}

// Keep track of current and previous steps,
// limiting them to the total number of steps
function useSteps(totalSteps?: number) {
  const [{ currentStep, previousStep }, dispatch] = useReducer(
    stepReducer,
    totalSteps,
    createInitialState,
  )
  const set = (step: number) => dispatch({ type: "set", payload: step })
  return { currentStep, previousStep, set }
}

type ActionType = { type: "set"; payload: number }
type StateType = ReturnType<typeof createInitialState>
function stepReducer(state: StateType, action: ActionType) {
  switch (action.type) {
    case "set": {
      if (action.payload === state.currentStep) return state
      if (action.payload < 0) return state
      if (state.totalSteps && action.payload > state.totalSteps - 1) return state
      return {
        ...state,
        currentStep: action.payload,
        previousStep: state.currentStep,
      }
    }
    default:
      throw new Error()
  }
}

function createInitialState(totalSteps?: number) {
  return { currentStep: 0, previousStep: 0, totalSteps }
}