Segmented control

"use client"

import * as React from "react"
import { cva, type VariantProps } from "class-variance-authority"

import { cn } from "@/lib/utils"
import { useSlidingIndicator } from "@/hooks/use-sliding-indicator"

const segmentedControlVariants = cva(
  "relative inline-flex items-center gap-0.5 rounded-lg bg-muted p-0.5 text-muted-foreground",
  {
    variants: {
      size: {
        default: "h-8",
        sm: "h-7",
      },
      fullWidth: {
        true: "flex w-full",
        false: "w-fit",
      },
    },
    defaultVariants: { size: "default", fullWidth: false },
  }
)

const segmentedControlOptionVariants = cva(
  // z-10, and no background of its own: the sliding pill behind the row is what
  // marks the selection, so a segment only owes the reader its ink.
  "relative z-10 inline-flex h-full min-w-0 items-center justify-center gap-1.5 rounded-[calc(var(--radius)-2px)] px-2.5 font-medium whitespace-nowrap transition-colors duration-(--duration-fast) ease-(--ease-standard) select-none focus-ring disabled:pointer-events-none disabled:opacity-50 [&_svg]:pointer-events-none [&_svg]:shrink-0 [&_svg:not([class*='size-'])]:size-4",
  {
    variants: {
      // The lift is the whole idea: the chosen segment sits on the sheet plane
      // while the track stays recessed.
      selected: {
        true: "text-foreground",
        false: "hover:text-foreground",
      },
      size: {
        default: "type-label",
        sm: "px-2 text-2xs [&_svg:not([class*='size-'])]:size-3.5",
      },
      fullWidth: {
        true: "flex-1",
        false: "",
      },
    },
    defaultVariants: { selected: false, size: "default", fullWidth: false },
  }
)

export type SegmentedControlOption = {
  value: string
  label: React.ReactNode
  /** Sits before the label; sized to 4 unless it sets its own size. */
  icon?: React.ReactNode
  disabled?: boolean
}

export type SegmentedControlProps = React.ComponentProps<"div"> & {
  options: SegmentedControlOption[]
  value: string
  onValueChange: (value: string) => void
  size?: NonNullable<VariantProps<typeof segmentedControlVariants>["size"]>
  /** Stretches the track and shares its width evenly between the segments. */
  fullWidth?: boolean
  "aria-label"?: string
}

/** A row of mutually exclusive options in one recessed track, for switching a view in place. */
function SegmentedControl({
  className,
  options,
  value,
  onValueChange,
  size = "default",
  fullWidth = false,
  ...props
}: SegmentedControlProps) {
  const optionRefs = React.useRef(new Map<string, HTMLButtonElement | null>())
  const { listRef, itemRef, style } = useSlidingIndicator({ value })

  // The first enabled option stands in for the tab stop when `value` matches
  // nothing, so the control is always reachable by keyboard.
  const selectedIndex = options.findIndex((option) => option.value === value)
  const tabStop =
    selectedIndex >= 0 && !options[selectedIndex].disabled
      ? options[selectedIndex].value
      : options.find((option) => !option.disabled)?.value

  // Walks outward from `from`, wrapping at either end, until it lands on an
  // option that can take the selection.
  function move(from: number, direction: 1 | -1) {
    const count = options.length
    for (let step = 1; step <= count; step += 1) {
      const next = options[(((from + direction * step) % count) + count) % count]
      if (next && !next.disabled) {
        onValueChange(next.value)
        optionRefs.current.get(next.value)?.focus()
        return
      }
    }
  }

  function handleKeyDown(event: React.KeyboardEvent<HTMLButtonElement>, index: number) {
    if (event.key === "ArrowRight" || event.key === "ArrowDown") {
      event.preventDefault()
      move(index, 1)
    } else if (event.key === "ArrowLeft" || event.key === "ArrowUp") {
      event.preventDefault()
      move(index, -1)
    }
  }

  return (
    <div
      data-slot="segmented-control"
      data-size={size}
      data-full-width={fullWidth || undefined}
      role="radiogroup"
      ref={listRef as React.RefObject<HTMLDivElement>}
      className={cn(segmentedControlVariants({ size, fullWidth }), className)}
      {...props}
    >
      <span
        data-slot="segmented-control-indicator"
        aria-hidden="true"
        style={style}
        className="absolute inset-y-0.5 left-0 z-0 rounded-[calc(var(--radius)-2px)] bg-card shadow-(--shadow-sheet) transition-[transform,width,opacity] duration-(--duration-base) ease-(--ease-emphasized)"
      />
      {options.map((option, index) => {
        const selected = option.value === value
        return (
          <button
            key={option.value}
            ref={(node) => {
              optionRefs.current.set(option.value, node)
              itemRef(selected)(node)
            }}
            type="button"
            role="radio"
            aria-checked={selected}
            data-slot="segmented-control-option"
            disabled={option.disabled}
            tabIndex={option.value === tabStop ? 0 : -1}
            onClick={() => onValueChange(option.value)}
            onKeyDown={(event) => handleKeyDown(event, index)}
            className={segmentedControlOptionVariants({ selected, size, fullWidth })}
          >
            {option.icon ? <span aria-hidden="true">{option.icon}</span> : null}
            <span className="truncate">{option.label}</span>
          </button>
        )
      })}
    </div>
  )
}

export { SegmentedControl, segmentedControlOptionVariants, segmentedControlVariants }