Toggle group

"use client";

import { createContext, useCallback, useContext, useEffect, useRef, useState } from "react";
import type { ComponentProps, FocusEvent, KeyboardEvent, MouseEvent, Ref } from "react";
import "./toggle-group.css";

type GroupContext = {
  isPressed: (value: string) => boolean;
  toggle: (value: string) => void;
  tabStop: string | null;
  setTabStop: (value: string) => void;
  disabled: boolean;
};

const ToggleGroupContext = createContext<GroupContext | null>(null);

export type ToggleGroupProps = Omit<ComponentProps<"div">, "defaultValue"> & {
  /** "single" keeps at most one item pressed; "multiple" lets each item toggle on its own. */
  type?: "single" | "multiple";
  /** Controlled list of pressed item values. */
  value?: string[];
  /** Uncontrolled starting list of pressed item values. */
  defaultValue?: string[];
  onValueChange?: (value: string[]) => void;
  /** Disables every item. */
  disabled?: boolean;
  /** What a row too long for its space does: scroll sideways, or wrap onto more lines. Unset, it overflows. */
  overflow?: "scroll" | "wrap";
  /** Height of the group: 28, 34 or 42px, to sit beside a Button of the same size. */
  size?: "s" | "m" | "l";
};

const ITEM = ".eq-toggle-group-item";

function items(root: HTMLElement | null) {
  return root ? Array.from(root.querySelectorAll<HTMLButtonElement>(ITEM)).filter((b) => !b.disabled) : [];
}

function setRefs<T>(node: T, ...refs: (Ref<T> | undefined)[]) {
  for (const ref of refs) {
    if (typeof ref === "function") ref(node);
    else if (ref) ref.current = node;
  }
}

export function ToggleGroup({
  type = "multiple",
  value,
  defaultValue = [],
  onValueChange,
  disabled = false,
  overflow,
  size = "m",
  className,
  onKeyDown,
  ref,
  children,
  ...rest
}: ToggleGroupProps) {
  const rootRef = useRef<HTMLDivElement | null>(null);
  const [inner, setInner] = useState(defaultValue);
  const [tabStop, setTabStop] = useState<string | null>(null);
  const pressed = value ?? inner;
  const mergedRef = useCallback((node: HTMLDivElement | null) => setRefs(node, rootRef, ref), [ref]);

  // Keep exactly one item in the tab order: the last one focused if it is still there,
  // else the first pressed item, else the first enabled item.
  useEffect(() => {
    const enabled = items(rootRef.current);
    if (enabled.some((b) => b.dataset.value === tabStop)) return;
    const next = enabled.find((b) => b.getAttribute("aria-pressed") === "true") ?? enabled[0];
    setTabStop(next?.dataset.value ?? null);
  });

  function toggle(item: string) {
    const on = pressed.includes(item);
    const next =
      type === "single" ? (on ? [] : [item]) : on ? pressed.filter((v) => v !== item) : [...pressed, item];
    if (value === undefined) setInner(next);
    onValueChange?.(next);
  }

  function handleKeyDown(event: KeyboardEvent<HTMLDivElement>) {
    onKeyDown?.(event);
    if (event.defaultPrevented) return;
    const enabled = items(rootRef.current);
    const index = enabled.findIndex((b) => b === document.activeElement);
    if (index === -1) return;
    let target: HTMLButtonElement | undefined;
    if (event.key === "ArrowRight") target = enabled[(index + 1) % enabled.length];
    else if (event.key === "ArrowLeft") target = enabled[(index - 1 + enabled.length) % enabled.length];
    else if (event.key === "Home") target = enabled[0];
    else if (event.key === "End") target = enabled[enabled.length - 1];
    if (!target) return;
    event.preventDefault();
    target.focus();
  }

  return (
    <ToggleGroupContext.Provider
      value={{ isPressed: (v) => pressed.includes(v), toggle, tabStop, setTabStop, disabled }}
    >
      <div
        role="toolbar"
        aria-orientation="horizontal"
        {...rest}
        ref={mergedRef}
        className={["eq-toggle-group", className].filter(Boolean).join(" ")}
        aria-disabled={disabled || undefined}
        data-type={type}
        data-overflow={overflow}
        data-size={size}
        onKeyDown={handleKeyDown}
      >
        {children}
      </div>
    </ToggleGroupContext.Provider>
  );
}

export type ToggleGroupItemProps = Omit<ComponentProps<"button">, "value"> & {
  /** Identifies the item in the group's value. */
  value: string;
};

export function ToggleGroupItem({
  value,
  className,
  disabled,
  onClick,
  onFocus,
  children,
  ...rest
}: ToggleGroupItemProps) {
  const group = useContext(ToggleGroupContext);
  if (!group) throw new Error("ToggleGroupItem must be used inside a ToggleGroup.");
  const isDisabled = disabled || group.disabled;
  // Before the group has picked a tab stop (on the server), every item stays reachable.
  const tabIndex = group.tabStop === null || group.tabStop === value ? 0 : -1;

  function handleClick(event: MouseEvent<HTMLButtonElement>) {
    onClick?.(event);
    if (!event.defaultPrevented) group?.toggle(value);
  }

  function handleFocus(event: FocusEvent<HTMLButtonElement>) {
    onFocus?.(event);
    group?.setTabStop(value);
  }

  return (
    <button
      type="button"
      {...rest}
      className={["eq-toggle-group-item", className].filter(Boolean).join(" ")}
      data-value={value}
      aria-pressed={group.isPressed(value)}
      disabled={isDisabled}
      tabIndex={tabIndex}
      onClick={handleClick}
      onFocus={handleFocus}
    >
      {children}
    </button>
  );
}