Segmented Control

"use client"

import * as React from "react"
import { motion } from "motion/react"

import { cn } from "@/lib/utils"
import { useControllableState } from "@/hooks/use-controllable-state"
import { spring } from "@/lib/motion"
import { useCurveTransition, type Curve } from "@/lib/transition"

export type SegmentedOption<T extends string = string> = {
  value: T
  label: React.ReactNode
  disabled?: boolean
}

export type SegmentedControlProps<T extends string = string> = Omit<
  React.ComponentProps<"div">,
  "onChange" | "defaultValue"
> & {
  /** The segments, in order. Pass an array of arrays to split them into sections with a divider between. */
  options: SegmentedOption<T>[] | SegmentedOption<T>[][]
  value?: T
  defaultValue?: T
  onValueChange?: (value: T) => void
  size?: "sm" | "md" | "lg"
  /** Stretch to the container and give every segment equal width. */
  fullWidth?: boolean
  disabled?: boolean
  /** Submits the selected value with a form. */
  name?: string
  /** Space between the track's edge and the thumb, in px. */
  inset?: number
  /** Space between segments, in px. */
  gap?: number
  /** Corner radius of the track, in px or any CSS length. The thumb follows it, less the inset. */
  radius?: number | string
  /** A hairline between every pair of segments, hidden beside the selected one. */
  dividers?: boolean
  /** How the thumb moves: any curve from the Animation Lab. A snappy spring when not set. */
  curve?: Curve
  /** Milliseconds for curves that are not springs. */
  duration?: number
  /** Extra classes for the parts. */
  classNames?: { thumb?: string; segment?: string; divider?: string }
}

const heights = { sm: "h-7 px-2.5 text-xs", md: "h-8 px-3 text-sm", lg: "h-10 px-4 text-sm" }
const length = (value: number | string) => (typeof value === "number" ? `${value}px` : value)

type ThumbPlacement = {
  value: string
  x: number
  y: number
  width: number
  height: number
  animate: boolean
}

function SegmentedControl<T extends string = string>({
  options,
  value: valueProp,
  defaultValue,
  onValueChange,
  size = "md",
  fullWidth = false,
  disabled = false,
  name,
  inset = 2,
  gap = 2,
  radius = 8,
  dividers = false,
  curve,
  duration,
  classNames,
  className,
  style,
  ...props
}: SegmentedControlProps<T>) {
  const sections = (Array.isArray(options[0]) ? options : [options]) as SegmentedOption<T>[][]
  const flat = sections.flat()
  const [value, setValue] = useControllableState<T | undefined>({
    value: valueProp,
    defaultValue: defaultValue ?? flat.find((o) => !o.disabled)?.value,
    onChange: onValueChange as (value: T | undefined) => void,
  })
  const refs = React.useRef(new Map<string, HTMLButtonElement>())
  const [thumb, setThumb] = React.useState<ThumbPlacement | null>(null)
  const transition = useCurveTransition(curve, duration, spring.hover)

  const measure = React.useCallback(() => {
    const selected = value === undefined ? undefined : refs.current.get(value)
    if (!selected || value === undefined) {
      setThumb(null)
      return
    }
    // Local offsets follow the track through scrolling, reflow and ancestor transforms.
    // Page coordinates would make the thumb animate the track's own movement again.
    const placement = {
      value,
      x: selected.offsetLeft,
      y: selected.offsetTop,
      width: selected.offsetWidth,
      height: selected.offsetHeight,
    }
    setThumb(current => {
      if (
        current && current.value === placement.value &&
        current.x === placement.x && current.y === placement.y &&
        current.width === placement.width && current.height === placement.height
      ) return current
      return { ...placement, animate: current !== null && current.value !== placement.value }
    })
  }, [value])

  // Measure after any prop/selection change, before painting. Resizing the track or its
  // labels repositions the thumb immediately; only a change of selection should slide.
  React.useLayoutEffect(() => { measure() })
  React.useLayoutEffect(() => {
    const observer = new ResizeObserver(measure)
    const track = refs.current.values().next().value?.parentElement
    if (track) observer.observe(track)
    refs.current.forEach(node => observer.observe(node))
    return () => observer.disconnect()
  }, [measure, options])

  const enabled = flat.filter((o) => !o.disabled && !disabled)
  // The tab stop is the selected segment, or the first enabled one.
  const tabStop = enabled.find((o) => o.value === value)?.value ?? enabled[0]?.value

  function onKeyDown(event: React.KeyboardEvent) {
    const index = enabled.findIndex((o) => o.value === tabStop)
    let next = index
    if (event.key === "ArrowRight" || event.key === "ArrowDown") next = index + 1
    else if (event.key === "ArrowLeft" || event.key === "ArrowUp") next = index - 1
    else if (event.key === "Home") next = 0
    else if (event.key === "End") next = enabled.length - 1
    else return
    event.preventDefault()
    const option = enabled[(next + enabled.length) % enabled.length]
    if (!option) return
    setValue(option.value)
    refs.current.get(option.value)?.focus()
  }

  // The thumb sits inside the track's padding, so its corners are the track's less the inset.
  const vars = {
    "--segmented-inset": `${inset}px`,
    "--segmented-gap": `${gap}px`,
    "--segmented-radius": length(radius),
    "--segmented-thumb-radius": `max(0px, calc(var(--segmented-radius) - var(--segmented-inset)))`,
  } as React.CSSProperties

  return (
    <div
      data-slot="segmented-control"
      role="radiogroup"
      aria-disabled={disabled || undefined}
      onKeyDown={onKeyDown}
      style={{ ...vars, ...style }}
      className={cn(
        "relative isolate inline-flex max-w-full items-stretch gap-(--segmented-gap) rounded-(--segmented-radius) bg-muted p-(--segmented-inset)",
        fullWidth && "flex w-full",
        disabled && "opacity-50",
        className
      )}
      {...props}
    >
      {thumb && (
        <motion.span
          aria-hidden="true"
          data-slot="segmented-control-thumb"
          initial={false}
          animate={{ x: thumb.x, y: thumb.y, width: thumb.width, height: thumb.height }}
          transition={thumb.animate ? transition : { duration: 0 }}
          className={cn(
            "pointer-events-none absolute top-0 left-0 z-0 rounded-(--segmented-thumb-radius) bg-background shadow-sm ring-1 ring-border",
            classNames?.thumb
          )}
        />
      )}
      {sections.map((section, s) => (
        <React.Fragment key={s}>
          {s > 0 && (
            <span
              aria-hidden="true"
              data-slot="segmented-control-divider"
              className={cn("mx-0.5 my-1.5 w-px shrink-0 self-stretch bg-foreground/15", classNames?.divider)}
            />
          )}
          {section.map((option, i) => {
            const checked = option.value === value
            const isDisabled = disabled || option.disabled
            // A divider between two segments hides when either side holds the thumb.
            const before = i > 0 ? section[i - 1] : undefined
            const hairline = dividers && before !== undefined
            return (
              <React.Fragment key={option.value}>
                {hairline && (
                  <span
                    aria-hidden="true"
                    data-slot="segmented-control-divider"
                    className={cn(
                      "-mx-[calc(var(--segmented-gap)/2_+_0.5px)] my-2 w-px shrink-0 self-stretch bg-foreground/12 transition-opacity duration-150",
                      (checked || before.value === value) && "opacity-0",
                      classNames?.divider
                    )}
                  />
                )}
                <button
                  ref={(node) => {
                    if (node) refs.current.set(option.value, node)
                    else refs.current.delete(option.value)
                  }}
                  type="button"
                  role="radio"
                  aria-checked={checked}
                  disabled={isDisabled}
                  tabIndex={option.value === tabStop ? 0 : -1}
                  onClick={() => setValue(option.value)}
                  className={cn(
                    "relative z-10 inline-flex min-w-0 cursor-pointer items-center justify-center rounded-(--segmented-thumb-radius) font-medium text-muted-foreground outline-none select-none",
                    "transition-colors duration-150 hover:text-foreground aria-checked:text-foreground",
                    "focus-visible:ring-2 focus-visible:ring-ring",
                    "disabled:pointer-events-none disabled:opacity-50",
                    heights[size],
                    fullWidth && "flex-1",
                    classNames?.segment
                  )}
                >
                  {checked && !thumb && (
                    <span
                      aria-hidden="true"
                      data-slot="segmented-control-thumb"
                      className={cn(
                        "pointer-events-none absolute inset-0 rounded-(--segmented-thumb-radius) bg-background shadow-sm ring-1 ring-border",
                        classNames?.thumb
                      )}
                    />
                  )}
                  <span className="relative inline-flex items-center gap-1.5 truncate [&_svg]:size-4 [&_svg]:shrink-0">
                    {option.label}
                  </span>
                </button>
              </React.Fragment>
            )
          })}
        </React.Fragment>
      ))}
      {name && value !== undefined && <input type="hidden" name={name} value={value} />}
    </div>
  )
}

export { SegmentedControl }