Badge Overflow

/**
 * @see https://github.com/radix-ui/primitives/blob/main/packages/react/compose-refs/src/compose-refs.tsx
 */
 
import * as React from "react";
 
type PossibleRef<T> = React.Ref<T> | undefined;
 
/**
 * Set a given ref to a given value
 * This utility takes care of different types of refs: callback refs and RefObject(s)
 */
function setRef<T>(ref: PossibleRef<T>, value: T) {
  if (typeof ref === "function") {
    return ref(value);
  }
 
  if (ref !== null && ref !== undefined) {
    ref.current = value;
  }
}
 
/**
 * A utility to compose multiple refs together
 * Accepts callback refs and RefObject(s)
 */
function composeRefs<T>(...refs: PossibleRef<T>[]): React.RefCallback<T> {
  return (node) => {
    let hasCleanup = false;
    const cleanups = refs.map((ref) => {
      const cleanup = setRef(ref, node);
      if (!hasCleanup && typeof cleanup === "function") {
        hasCleanup = true;
      }
      return cleanup;
    });
 
    // React <19 will log an error to the console if a callback ref returns a
    // value. We don't use ref cleanups internally so this will only happen if a
    // user's ref callback returns a value, which we only expect if they are
    // using the cleanup functionality added in React 19.
    if (hasCleanup) {
      return () => {
        for (let i = 0; i < cleanups.length; i++) {
          const cleanup = cleanups[i];
          if (typeof cleanup === "function") {
            cleanup();
          } else {
            setRef(refs[i], null);
          }
        }
      };
    }
  };
}
 
/**
 * A custom hook that composes multiple refs
 * Accepts callback refs and RefObject(s)
 */
function useComposedRefs<T>(...refs: PossibleRef<T>[]): React.RefCallback<T> {
  // biome-ignore lint/correctness/useExhaustiveDependencies: we want to memoize by all values
  return React.useCallback(composeRefs(...refs), refs);
}
 
export { composeRefs, useComposedRefs };

"use client";
 
import { Slot } from "@radix-ui/react-slot";
import * as React from "react";
import { useComposedRefs } from "@/lib/compose-refs";
import { cn } from "@/lib/utils";
 
interface GetBadgeLabel<T> {
  /**
   * Callback that returns a label string for each badge item.
   * Optional for primitive arrays (strings, numbers), required for object arrays.
   * @example getBadgeLabel={(item) => item.name}
   */
  getBadgeLabel: (item: T) => string;
}
 
type BadgeOverflowElement = React.ComponentRef<typeof BadgeOverflow>;
 
type BadgeOverflowProps<T = string> = React.ComponentProps<"div"> &
  (T extends object ? GetBadgeLabel<T> : Partial<GetBadgeLabel<T>>) & {
    items: T[];
    lineCount?: number;
    renderBadge: (item: T, label: string) => React.ReactNode;
    renderOverflow?: (count: number) => React.ReactNode;
    asChild?: boolean;
  };
 
function BadgeOverflow<T = string>(props: BadgeOverflowProps<T>) {
  const {
    items,
    getBadgeLabel: getBadgeLabelProp,
    lineCount = 1,
    renderBadge,
    renderOverflow,
    asChild,
    className,
    style,
    ref,
    ...rootProps
  } = props;
 
  const getBadgeLabel = React.useCallback(
    (item: T): string => {
      if (typeof item === "object" && !getBadgeLabelProp) {
        throw new Error(
          "`getBadgeLabel` is required when using array of objects",
        );
      }
      return getBadgeLabelProp ? getBadgeLabelProp(item) : (item as string);
    },
    [getBadgeLabelProp],
  );
 
  const rootRef = React.useRef<BadgeOverflowElement | null>(null);
  const composedRef = useComposedRefs(ref, rootRef);
  const measureRef = React.useRef<HTMLDivElement | null>(null);
  const [containerWidth, setContainerWidth] = React.useState(0);
  const [badgeGap, setBadgeGap] = React.useState(4);
  const [badgeHeight, setBadgeHeight] = React.useState(20);
  const [overflowBadgeWidth, setOverflowBadgeWidth] = React.useState(40);
  const [isMeasured, setIsMeasured] = React.useState(false);
  const [badgeWidths, setBadgeWidths] = React.useState<Map<string, number>>(
    new Map(),
  );
 
  React.useLayoutEffect(() => {
    if (!rootRef.current || !measureRef.current) return;
 
    function measureContainer() {
      if (!rootRef.current || !measureRef.current) return;
 
      const computedStyle = getComputedStyle(rootRef.current);
 
      const gapValue = computedStyle.gap;
      const gap = gapValue ? parseFloat(gapValue) : 4;
      setBadgeGap(gap);
 
      const paddingLeft = parseFloat(computedStyle.paddingLeft) || 0;
      const paddingRight = parseFloat(computedStyle.paddingRight) || 0;
      const totalPadding = paddingLeft + paddingRight;
 
      const widthMap = new Map<string, number>();
      const measureChildren = measureRef.current.children;
 
      for (let i = 0; i < items.length; i++) {
        const child = measureChildren[i] as HTMLElement | undefined;
        if (child) {
          const label = getBadgeLabel(items[i] as T);
          widthMap.set(label, child.offsetWidth);
        }
      }
      setBadgeWidths(widthMap);
 
      const firstBadge = measureChildren[0] as HTMLElement | undefined;
      if (firstBadge) {
        setBadgeHeight(firstBadge.offsetHeight || 20);
      }
 
      const overflowChild = measureChildren[items.length] as
        | HTMLElement
        | undefined;
 
      if (overflowChild) {
        setOverflowBadgeWidth(overflowChild.offsetWidth || 40);
      }
 
      const width = rootRef.current.clientWidth - totalPadding;
      setContainerWidth(width);
      setIsMeasured(true);
    }
 
    measureContainer();
 
    const resizeObserver = new ResizeObserver(measureContainer);
    resizeObserver.observe(rootRef.current);
 
    return () => {
      resizeObserver.disconnect();
    };
  }, [items, getBadgeLabel]);
 
  const placeholderHeight = React.useMemo(
    () => badgeHeight * lineCount + badgeGap * (lineCount - 1),
    [badgeHeight, badgeGap, lineCount],
  );
 
  const { visibleItems, hiddenCount } = React.useMemo(() => {
    if (!containerWidth || items.length === 0 || badgeWidths.size === 0) {
      return { visibleItems: items, hiddenCount: 0 };
    }
 
    let currentLineWidth = 0;
    let currentLine = 1;
    const visible: T[] = [];
 
    for (let i = 0; i < items.length; i++) {
      const item = items[i];
      if (!item) continue;
 
      const label = getBadgeLabel(item);
      const badgeWidth = badgeWidths.get(label);
 
      if (!badgeWidth) {
        // Skip items that haven't been measured yet
        continue;
      }
 
      const widthWithGap = badgeWidth + badgeGap;
      const isLastLine = currentLine === lineCount;
      const hasMoreItems = i < items.length - 1;
 
      const availableWidth =
        isLastLine && hasMoreItems
          ? containerWidth - overflowBadgeWidth - badgeGap
          : containerWidth;
 
      if (currentLineWidth + widthWithGap <= availableWidth) {
        currentLineWidth += widthWithGap;
        visible.push(item);
      } else if (currentLine < lineCount) {
        currentLine++;
        currentLineWidth = widthWithGap;
        visible.push(item);
      } else {
        // We're on the last line and this badge doesn't fit
        break;
      }
    }
 
    return {
      visibleItems: visible,
      hiddenCount: Math.max(0, items.length - visible.length),
    };
  }, [
    items,
    getBadgeLabel,
    containerWidth,
    lineCount,
    badgeGap,
    overflowBadgeWidth,
    badgeWidths,
  ]);
 
  const Comp = asChild ? Slot : "div";
 
  return (
    <>
      <div
        ref={measureRef}
        className="pointer-events-none invisible absolute flex flex-wrap"
        style={{ gap: badgeGap }}
      >
        {items.map((item, index) => (
          <React.Fragment key={index}>
            {renderBadge(item, getBadgeLabel(item))}
          </React.Fragment>
        ))}
        {renderOverflow ? (
          renderOverflow(99)
        ) : (
          <div className="inline-flex h-5 shrink-0 items-center rounded-md border px-1.5 font-semibold text-xs">
            +99
          </div>
        )}
      </div>
      {isMeasured ? (
        <Comp
          data-slot="badge-overflow"
          {...rootProps}
          ref={composedRef}
          className={cn("flex flex-wrap", className)}
          style={{
            gap: badgeGap,
            ...style,
          }}
        >
          {visibleItems.map((item, index) => (
            <React.Fragment key={index}>
              {renderBadge(item, getBadgeLabel(item))}
            </React.Fragment>
          ))}
          {hiddenCount > 0 &&
            (renderOverflow ? (
              renderOverflow(hiddenCount)
            ) : (
              <div className="inline-flex h-5 shrink-0 items-center rounded-md border px-1.5 font-semibold text-xs">
                +{hiddenCount}
              </div>
            ))}
        </Comp>
      ) : (
        <Comp
          data-slot="badge-overflow"
          {...rootProps}
          ref={composedRef}
          className={cn("flex flex-wrap", className)}
          style={{
            gap: badgeGap,
            minHeight: placeholderHeight,
            ...style,
          }}
        >
          {items
            .slice(
              0,
              Math.min(items.length, lineCount * 3 - (lineCount > 1 ? 1 : 0)),
            )
            .map((item, index) => (
              <React.Fragment key={index}>
                {renderBadge(item, getBadgeLabel(item))}
              </React.Fragment>
            ))}
        </Comp>
      )}
    </>
  );
}
 
export { BadgeOverflow };