import { createContext, forwardRef, useContext } from 'react'
import type { ReactNode } from 'react'
import * as TogglePrimitive from '@radix-ui/react-toggle'
import * as ToggleGroupPrimitive from '@radix-ui/react-toggle-group'
import { cva, type VariantProps } from 'class-variance-authority'
import { cn } from '@/lib/utils'

const toggleVariants = cva(
  'inline-flex items-center justify-center gap-1.5 font-medium transition-colors duration-200 focus:outline-none focus-visible:ring-2 focus-visible:ring-primary-border focus-visible:ring-offset-2 focus-visible:ring-offset-bg-base disabled:pointer-events-none disabled:opacity-50 cursor-pointer data-[state=on]:bg-fill data-[state=on]:text-text',
  {
    variants: {
      variant: {
        default:
          'bg-transparent text-text-secondary hover:bg-fill-secondary hover:text-text data-[state=on]:bg-fill data-[state=on]:text-text',
        outline:
          'border border-border-tertiary bg-transparent text-text-secondary hover:bg-fill-secondary hover:text-text data-[state=on]:border-border-secondary data-[state=on]:bg-fill data-[state=on]:text-text',
      },
      size: {
        sm: 'h-[var(--control-height-sm)] px-[var(--control-padding-x-sm)] text-[length:var(--control-font-size-sm)] leading-[var(--control-line-height-sm)]',
        md: 'h-[var(--control-height-md)] px-[var(--control-padding-x-md)] text-[length:var(--control-font-size-md)] leading-[var(--control-line-height-md)]',
        lg: 'h-[var(--control-height-lg)] px-[var(--control-padding-x-lg)] text-[length:var(--control-font-size-lg)] leading-[var(--control-line-height-lg)]',
      },
      rounded: {
        square: 'rounded',
        pill: 'rounded-full',
      },
    },
    defaultVariants: { variant: 'default', size: 'md', rounded: 'square' },
  }
)

export interface ToggleProps
  extends Omit<React.ComponentPropsWithoutRef<typeof TogglePrimitive.Root>, 'asChild'>,
    VariantProps<typeof toggleVariants> {
  children?: ReactNode
  prefixIcon?: ReactNode
  suffixIcon?: ReactNode
}

export const Toggle = forwardRef<
  React.ElementRef<typeof TogglePrimitive.Root>,
  ToggleProps
>(
  (
    {
      variant = 'default',
      size = 'md',
      rounded = 'square',
      className,
      prefixIcon,
      suffixIcon,
      children,
      ...props
    },
    ref
  ) => (
    <TogglePrimitive.Root
      ref={ref}
      className={toggleVariants({ variant, size, rounded, className })}
      {...props}
    >
      {prefixIcon && (
        <span className="inline-flex shrink-0 items-center justify-center [&>*]:block [&>*]:leading-none">
          {prefixIcon}
        </span>
      )}
      {children}
      {suffixIcon && (
        <span className="inline-flex shrink-0 items-center justify-center [&>*]:block [&>*]:leading-none">
          {suffixIcon}
        </span>
      )}
    </TogglePrimitive.Root>
  )
)
Toggle.displayName = 'Toggle'

type ToggleGroupContextValue = VariantProps<typeof toggleVariants> & {
  spacing?: 'none' | 'sm' | 'md'
}

const ToggleGroupContext = createContext<ToggleGroupContextValue | null>(null)
const spacingClasses = { none: 'gap-0', sm: 'gap-1', md: 'gap-2' } as const

const toggleGroupItemGroupClasses = [
  'shrink-0',
  'group-data-[spacing=0]/toggle-group:rounded-none',
  'group-data-[spacing=0]/toggle-group:group-data-[orientation=horizontal]/toggle-group:first:rounded-l',
  'group-data-[spacing=0]/toggle-group:group-data-[orientation=horizontal]/toggle-group:last:rounded-r',
  'group-data-[spacing=0]/toggle-group:group-data-[orientation=vertical]/toggle-group:first:rounded-t',
  'group-data-[spacing=0]/toggle-group:group-data-[orientation=vertical]/toggle-group:last:rounded-b',
  'group-data-[spacing=0]/toggle-group:data-[variant=outline]:group-data-[orientation=horizontal]/toggle-group:[&:not(:first-child)]:border-l-0',
  'group-data-[spacing=0]/toggle-group:data-[variant=outline]:group-data-[orientation=vertical]/toggle-group:[&:not(:first-child)]:border-t-0',
].join(' ')

type ToggleGroupRootProps = React.ComponentPropsWithoutRef<typeof ToggleGroupPrimitive.Root>

export type ToggleGroupProps = ToggleGroupRootProps &
  VariantProps<typeof toggleVariants> & {
    spacing?: 'none' | 'sm' | 'md'
    children?: ReactNode
  }

const ToggleGroup = forwardRef<
  React.ElementRef<typeof ToggleGroupPrimitive.Root>,
  ToggleGroupProps
>(
  (
    {
      type = 'single',
      variant = 'default',
      size = 'md',
      rounded = 'square',
      spacing = 'none',
      className,
      orientation = 'horizontal',
      children,
      ...props
    },
    ref
  ) => {
    const rootClassName = cn(
      'group/toggle-group inline-flex w-fit',
      orientation === 'horizontal' ? 'flex-row' : 'flex-col',
      orientation === 'vertical' && 'items-stretch',
      spacingClasses[spacing],
      className
    )
    const commonProps = {
      ref,
      orientation,
      'data-spacing': spacing === 'none' ? '0' : spacing,
      'data-orientation': orientation,
      'data-variant': variant,
      'data-size': size,
      className: rootClassName,
    } as const
    return (
      <ToggleGroupContext.Provider value={{ variant, size, rounded, spacing }}>
        {type === 'multiple' ? (
          <ToggleGroupPrimitive.Root
            {...commonProps}
            {...(props as ToggleGroupPrimitive.ToggleGroupMultipleProps)}
            type="multiple"
          >
            {children}
          </ToggleGroupPrimitive.Root>
        ) : (
          <ToggleGroupPrimitive.Root
            {...commonProps}
            {...(props as ToggleGroupPrimitive.ToggleGroupSingleProps)}
            type="single"
          >
            {children}
          </ToggleGroupPrimitive.Root>
        )}
      </ToggleGroupContext.Provider>
    )
  }
)
ToggleGroup.displayName = 'ToggleGroup'

export interface ToggleGroupItemProps
  extends Omit<React.ComponentPropsWithoutRef<typeof ToggleGroupPrimitive.Item>, 'asChild'>,
    VariantProps<typeof toggleVariants> {
  value: string
  children?: ReactNode
  prefixIcon?: ReactNode
  suffixIcon?: ReactNode
}

const ToggleGroupItem = forwardRef<
  React.ElementRef<typeof ToggleGroupPrimitive.Item>,
  ToggleGroupItemProps
>(
  (
    {
      variant: variantProp,
      size: sizeProp,
      rounded: roundedProp,
      className,
      prefixIcon,
      suffixIcon,
      children,
      ...props
    },
    ref
  ) => {
    const ctx = useContext(ToggleGroupContext)
    const variant = variantProp ?? ctx?.variant ?? 'default'
    const size = sizeProp ?? ctx?.size ?? 'md'
    const rounded = roundedProp ?? ctx?.rounded ?? 'square'
    return (
      <ToggleGroupPrimitive.Item
        ref={ref}
        data-variant={variant}
        data-size={size}
        className={cn(
          toggleVariants({ variant, size, rounded }),
          ctx && toggleGroupItemGroupClasses,
          'focus-visible:z-10',
          className
        )}
        {...props}
      >
        {prefixIcon && (
          <span className="inline-flex shrink-0 items-center justify-center [&>*]:block [&>*]:leading-none">
            {prefixIcon}
          </span>
        )}
        {children}
        {suffixIcon && (
          <span className="inline-flex shrink-0 items-center justify-center [&>*]:block [&>*]:leading-none">
            {suffixIcon}
          </span>
        )}
      </ToggleGroupPrimitive.Item>
    )
  }
)
ToggleGroupItem.displayName = 'ToggleGroupItem'

export { ToggleGroup, ToggleGroupItem, toggleVariants }
