'use client'

import * as ToggleGroupPrimitive from '@radix-ui/react-toggle-group'
import * as React from 'react'
import { cn } from '../../utils'
import { segmentButtonVariants, type VariantProps } from '../SegmentButton'

const SegmentControllerContext = React.createContext<{
  size?: VariantProps<typeof segmentButtonVariants>['size']
} | null>(null)

interface SegmentControllerProps
  extends Omit<
    React.ComponentPropsWithoutRef<typeof ToggleGroupPrimitive.Root>,
    'type' | 'value' | 'defaultValue' | 'onValueChange'
  > {
  className?: string
  children?: React.ReactNode
  value?: string
  defaultValue?: string
  onValueChange?: (value: string) => void
  size?: VariantProps<typeof segmentButtonVariants>['size']
  ref?: React.Ref<React.ElementRef<typeof ToggleGroupPrimitive.Root>>
}

const SegmentController = ({
  className,
  size,
  children,
  value,
  defaultValue,
  onValueChange,
  ref,
  ...props
}: SegmentControllerProps) => {
  const [internalValue, setInternalValue] = React.useState<string>(
    value ?? defaultValue ?? '',
  )
  const isControlled = value !== undefined
  const currentValue = isControlled ? value : internalValue

  const handleValueChange = React.useCallback(
    (newValue: string | undefined) => {
      if (!newValue) return

      if (!isControlled) {
        setInternalValue(newValue)
      }

      onValueChange?.(newValue)
    },
    [isControlled, onValueChange],
  )

  const contextValue = React.useMemo(() => ({ size }), [size])

  return (
    <SegmentControllerContext.Provider value={contextValue}>
      <ToggleGroupPrimitive.Root
        {...props}
        ref={ref}
        type="single"
        value={currentValue}
        onValueChange={handleValueChange}
        className={cn(
          'inline-flex items-center justify-center overflow-hidden rounded-md border border-solid border-segment-button-border',
          className,
        )}
      >
        {children}
      </ToggleGroupPrimitive.Root>
    </SegmentControllerContext.Provider>
  )
}

SegmentController.displayName = 'SegmentController'

type SegmentControllerItemProps = React.ComponentPropsWithoutRef<
  typeof ToggleGroupPrimitive.Item
> &
  VariantProps<typeof segmentButtonVariants> & {
    ref?: React.Ref<React.ElementRef<typeof ToggleGroupPrimitive.Item>>
  }

const SegmentControllerItem = ({
  className,
  children,
  variant,
  size,
  ref,
  ...props
}: SegmentControllerItemProps) => {
  const context = React.useContext(SegmentControllerContext)

  const itemVariant = variant ?? 'default'
  const itemSize = context?.size ?? size

  return (
    <ToggleGroupPrimitive.Item
      ref={ref}
      className={cn(
        segmentButtonVariants({ variant: itemVariant, size: itemSize }),
        'w-full rounded-none border-r border-segment-button-border last:border-r-0',
        'bg-segment-button-background text-segment-button-foreground',
        'hover:bg-segment-button-background-hover',
        'data-[state=on]:bg-segment-button-active data-[state=on]:text-segment-button-active-foreground',
        className,
      )}
      {...props}
    >
      {children}
    </ToggleGroupPrimitive.Item>
  )
}

SegmentControllerItem.displayName = 'SegmentControllerItem'

export { SegmentController, SegmentControllerItem }