Toggle

import type { VariantProps } from "cva";
import {
  Children,
  createContext,
  isValidElement,
  use,
  useCallback,
  useLayoutEffect,
  useMemo,
  useRef,
  useState,
} from "react";
import { Slot } from "@/components/slot";
import { composeRefs } from "@/lib/compose-refs";
import { cn } from "@/lib/utils/classnames";
import { buttonStyle } from "../button";

type ToggleStyleProps = Pick<VariantProps<typeof buttonStyle>, "size" | "square">;

interface ToggleProps
  extends Omit<React.ComponentPropsWithRef<"button">, "onChange">,
    ToggleStyleProps {
  pressed?: boolean;
  defaultPressed?: boolean;
  onPressedChange?: (pressed: boolean) => void;
  asChild?: boolean;
}

const Toggle = ({
  pressed: pressedProp,
  defaultPressed = false,
  onPressedChange,
  size = "md",
  square,
  asChild = false,
  className,
  type = "button",
  onClick,
  ref,
  children,
  ...props
}: ToggleProps) => {
  const [internalPressed, setInternalPressed] = useState(defaultPressed);
  const pressed = pressedProp ?? internalPressed;

  const Comp = asChild ? Slot : "button";

  return (
    <Comp
      ref={ref}
      type={asChild ? undefined : type}
      className={cn(buttonStyle({ variant: "ghost", size, square }), className)}
      aria-pressed={pressed}
      data-pressed={pressed || undefined}
      onClick={(e) => {
        onClick?.(e);
        if (e.defaultPrevented) return;
        const next = !pressed;
        if (pressedProp === undefined) setInternalPressed(next);
        onPressedChange?.(next);
      }}
      {...props}
    >
      {children}
    </Comp>
  );
};

interface ToggleGroupContextValue {
  isPressed: (value: string) => boolean;
  toggle: (value: string) => void;
  tabbableValue: string | undefined;
  setTabbableValue: (value: string) => void;
  focusItem: (value: string) => void;
  registerItem: (value: string) => () => void;
  items: string[];
  orientation: "horizontal" | "vertical";
  disabled: boolean;
  size: ToggleStyleProps["size"];
}

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

const useToggleGroupContext = () => {
  const ctx = use(ToggleGroupContext);
  if (!ctx) {
    throw new Error("ToggleGroup.Item must be used within a ToggleGroup");
  }
  return ctx;
};

type ToggleGroupBaseProps = Omit<React.ComponentPropsWithRef<"div">, "onChange" | "defaultValue"> &
  Pick<ToggleStyleProps, "size"> & {
    orientation?: "horizontal" | "vertical";
    disabled?: boolean;
  };

type ToggleGroupSingleProps = ToggleGroupBaseProps & {
  type: "single";
  value?: string;
  defaultValue?: string;
  onValueChange?: (value: string | undefined) => void;
};

type ToggleGroupMultipleProps = ToggleGroupBaseProps & {
  type: "multiple";
  value?: string[];
  defaultValue?: string[];
  onValueChange?: (value: string[]) => void;
};

type ToggleGroupProps = ToggleGroupSingleProps | ToggleGroupMultipleProps;

const ToggleGroup = (props: ToggleGroupProps) => {
  return props.type === "single" ? (
    <ToggleGroupSingle {...props} />
  ) : (
    <ToggleGroupMultiple {...props} />
  );
};

const ToggleGroupSingle = ({
  type: _type,
  value: valueProp,
  defaultValue,
  onValueChange,
  ...rest
}: ToggleGroupSingleProps) => {
  const [internal, setInternal] = useState<string | undefined>(defaultValue);
  const value = valueProp ?? internal;

  const isPressed = useCallback((v: string) => value === v, [value]);
  const toggle = useCallback(
    (v: string) => {
      const next = value === v ? undefined : v;
      if (valueProp === undefined) setInternal(next);
      onValueChange?.(next);
    },
    [value, valueProp, onValueChange],
  );

  return <ToggleGroupRoot {...rest} isPressed={isPressed} toggle={toggle} />;
};

const ToggleGroupMultiple = ({
  type: _type,
  value: valueProp,
  defaultValue,
  onValueChange,
  ...rest
}: ToggleGroupMultipleProps) => {
  const [internal, setInternal] = useState<string[]>(defaultValue ?? []);
  const value = valueProp ?? internal;

  const isPressed = useCallback((v: string) => value.includes(v), [value]);
  const toggle = useCallback(
    (v: string) => {
      const next = value.includes(v) ? value.filter((x) => x !== v) : [...value, v];
      if (valueProp === undefined) setInternal(next);
      onValueChange?.(next);
    },
    [value, valueProp, onValueChange],
  );

  return <ToggleGroupRoot {...rest} isPressed={isPressed} toggle={toggle} />;
};

interface ToggleGroupRootProps extends ToggleGroupBaseProps {
  isPressed: (value: string) => boolean;
  toggle: (value: string) => void;
}

const ToggleGroupRoot = ({
  isPressed,
  toggle,
  orientation = "horizontal",
  disabled = false,
  size = "md",
  className,
  children,
  ref,
  ...divProps
}: ToggleGroupRootProps) => {
  const [items, setItems] = useState<string[]>([]);

  const registerItem = useCallback((value: string) => {
    setItems((prev) => (prev.includes(value) ? prev : [...prev, value]));
    return () => {
      setItems((prev) => prev.filter((v) => v !== value));
    };
  }, []);

  const containerRef = useRef<HTMLDivElement>(null);

  const focusItem = useCallback((value: string) => {
    const el = containerRef.current?.querySelector<HTMLElement>(
      `[data-toggle-group-value="${CSS.escape(value)}"]`,
    );
    el?.focus();
  }, []);

  const [tabbable, setTabbable] = useState<string | undefined>(undefined);

  const tabbableValue = useMemo(() => {
    if (items.length > 0) {
      if (tabbable && items.includes(tabbable)) return tabbable;
      const firstPressed = items.find((v) => isPressed(v));
      return firstPressed ?? items[0];
    }
    // Pre-registration (SSR / first paint): items register via useLayoutEffect,
    // so on the server they're empty. Walk children directly so the right item
    // is the tab stop on the very first render.
    const childValues: string[] = [];
    Children.forEach(children, (child) => {
      if (isValidElement<{ value?: unknown }>(child) && typeof child.props.value === "string") {
        childValues.push(child.props.value);
      }
    });
    const firstPressed = childValues.find((v) => isPressed(v));
    return firstPressed ?? childValues[0];
  }, [tabbable, items, isPressed, children]);

  const setTabbableValue = useCallback((value: string) => {
    setTabbable(value);
  }, []);

  const ctx = useMemo<ToggleGroupContextValue>(
    () => ({
      isPressed,
      toggle,
      tabbableValue,
      setTabbableValue,
      focusItem,
      registerItem,
      items,
      orientation,
      disabled,
      size,
    }),
    [
      isPressed,
      toggle,
      tabbableValue,
      setTabbableValue,
      focusItem,
      registerItem,
      items,
      orientation,
      disabled,
      size,
    ],
  );

  return (
    <ToggleGroupContext value={ctx}>
      {/* biome-ignore lint/a11y/useSemanticElements: <fieldset> is form-only semantics; this is a toolbar-style grouping. */}
      <div
        ref={composeRefs(containerRef, ref)}
        role="group"
        data-orientation={orientation}
        className={cn(
          "flex items-center *:focus-visible:z-2",
          orientation === "vertical" && "flex-col",
          className,
        )}
        data-ui-button-group
        {...divProps}
      >
        {children}
      </div>
    </ToggleGroupContext>
  );
};

interface ToggleGroupItemProps
  extends Omit<React.ComponentPropsWithRef<"button">, "value" | "type" | "disabled" | "onChange">,
    ToggleStyleProps {
  value: string;
  asChild?: boolean;
  disabled?: boolean;
}

const ToggleGroupItem = ({
  value,
  size: sizeProp,
  square,
  asChild = false,
  className,
  disabled: disabledProp,
  onClick,
  onKeyDown,
  onFocus,
  ref,
  children,
  ...props
}: ToggleGroupItemProps) => {
  const ctx = useToggleGroupContext();
  const Comp = asChild ? Slot : "button";

  useLayoutEffect(() => {
    return ctx.registerItem(value);
  }, [value, ctx.registerItem]);

  const pressed = ctx.isPressed(value);
  const isTabbable = ctx.tabbableValue === value;
  const disabled = disabledProp || ctx.disabled;
  const size = sizeProp ?? ctx.size;

  const handleKeyDown = (e: React.KeyboardEvent<HTMLButtonElement>) => {
    const orientation = ctx.orientation;
    const nextKey = orientation === "horizontal" ? "ArrowRight" : "ArrowDown";
    const prevKey = orientation === "horizontal" ? "ArrowLeft" : "ArrowUp";

    if (e.key !== nextKey && e.key !== prevKey && e.key !== "Home" && e.key !== "End") {
      return;
    }

    e.preventDefault();
    const idx = ctx.items.indexOf(value);
    if (idx === -1) return;

    let nextIdx = idx;
    if (e.key === nextKey) nextIdx = (idx + 1) % ctx.items.length;
    else if (e.key === prevKey) nextIdx = (idx - 1 + ctx.items.length) % ctx.items.length;
    else if (e.key === "Home") nextIdx = 0;
    else if (e.key === "End") nextIdx = ctx.items.length - 1;

    const nextValue = ctx.items[nextIdx];
    if (!nextValue) return;

    ctx.setTabbableValue(nextValue);
    ctx.focusItem(nextValue);
  };

  return (
    <Comp
      ref={ref}
      type={asChild ? undefined : "button"}
      data-toggle-group-value={value}
      className={cn(buttonStyle({ variant: "ghost", size, square }), className)}
      aria-pressed={pressed}
      data-pressed={pressed || undefined}
      data-disabled={disabled || undefined}
      disabled={disabled}
      tabIndex={isTabbable ? 0 : -1}
      onClick={(e) => {
        onClick?.(e);
        if (e.defaultPrevented) return;
        ctx.toggle(value);
        ctx.setTabbableValue(value);
      }}
      onKeyDown={(e) => {
        onKeyDown?.(e);
        if (!e.defaultPrevented) handleKeyDown(e);
      }}
      onFocus={(e) => {
        onFocus?.(e);
        if (!e.defaultPrevented) ctx.setTabbableValue(value);
      }}
      {...props}
    >
      {children}
    </Comp>
  );
};

const CompoundToggleGroup = Object.assign(ToggleGroup, {
  Item: ToggleGroupItem,
});

export type { ToggleGroupItemProps, ToggleGroupProps, ToggleProps };
export { CompoundToggleGroup as ToggleGroup, Toggle };