import * as ToggleGroupPrimitive from "@radix-ui/react-toggle-group";
import { toggleVariants } from "@repo/shadcn-ui/components/toggle";
import { cn } from "@repo/shadcn-ui/lib/utils";
import type { VariantProps } from "class-variance-authority";
import * as React from "react";

const ToggleGroupContext = React.createContext<
  VariantProps<typeof toggleVariants> & {
    spacing?: number;
    orientation?: "horizontal" | "vertical";
  }
>({
  size: "default",
  variant: "default",
  spacing: 0,
  orientation: "horizontal",
});

const ToggleGroup = React.forwardRef<
  React.ElementRef<typeof ToggleGroupPrimitive.Root>,
  React.ComponentPropsWithoutRef<typeof ToggleGroupPrimitive.Root> &
    VariantProps<typeof toggleVariants> & {
      spacing?: number;
    }
>(({ className, variant, size, spacing = 0, children, style, ...props }, ref) => {
  const orientation = props.orientation ?? "horizontal";

  return (
    <ToggleGroupPrimitive.Root
      className={cn(
        "group/toggle-group flex w-fit rounded-md",
        orientation === "vertical" ? "flex-col items-stretch" : "flex-row items-center",
        className
      )}
      data-orientation={orientation}
      data-size={size}
      data-slot="toggle-group"
      data-spacing={spacing}
      data-variant={variant}
      ref={ref}
      style={{
        gap: spacing === 0 ? undefined : `${spacing * 0.25}rem`,
        ...style,
      }}
      {...props}
    >
      <ToggleGroupContext.Provider value={{ variant, size, spacing, orientation }}>
        {children}
      </ToggleGroupContext.Provider>
    </ToggleGroupPrimitive.Root>
  );
});
ToggleGroup.displayName = ToggleGroupPrimitive.Root.displayName;

const ToggleGroupItem = React.forwardRef<
  React.ElementRef<typeof ToggleGroupPrimitive.Item>,
  React.ComponentPropsWithoutRef<typeof ToggleGroupPrimitive.Item> &
    VariantProps<typeof toggleVariants>
>(({ className, children, variant, size, ...props }, ref) => {
  const context = React.useContext(ToggleGroupContext);

  return (
    <ToggleGroupPrimitive.Item
      className={cn(
        "shrink-0 rounded-md focus:z-10 focus-visible:z-10 data-[state=on]:bg-muted data-[state=on]:text-foreground",
        "group-data-[spacing=0]/toggle-group:first:rounded-r-none group-data-[spacing=0]/toggle-group:last:rounded-l-none group-data-[spacing=0]/toggle-group:not-first:not-last:rounded-none",
        "group-data-[spacing=0]/toggle-group:data-[orientation=vertical]:first:rounded-b-none group-data-[spacing=0]/toggle-group:data-[orientation=vertical]:last:rounded-t-none",
        "group-data-[spacing=0]/toggle-group:data-[orientation=vertical]:not-first:not-last:rounded-none",
        "group-data-[spacing=0]/toggle-group:not-first:-ml-px group-data-[spacing=0]/toggle-group:data-[orientation=vertical]:not-first:ml-0 group-data-[spacing=0]/toggle-group:data-[orientation=vertical]:not-first:-mt-px",
        toggleVariants({
          variant: context.variant ?? variant ?? "default",
          size: context.size ?? size ?? "default",
        }),
        className
      )}
      data-orientation={context.orientation}
      data-size={context.size || size}
      data-slot="toggle-group-item"
      data-spacing={context.spacing}
      data-variant={context.variant || variant}
      ref={ref}
      {...props}
    >
      {children}
    </ToggleGroupPrimitive.Item>
  );
});
ToggleGroupItem.displayName = ToggleGroupPrimitive.Item.displayName;

export { ToggleGroup, ToggleGroupItem };