import * as React from 'react'
import { cn } from '../lib/cn'

type GridGap = 2 | 4 | 6 | 8 | 12 | 16 | 24

export interface StackProps extends React.HTMLAttributes<HTMLDivElement> {
  direction?: 'vertical' | 'horizontal'
  /** Gap restricted to the 8px-grid tokens. */
  gap?: GridGap
  align?: 'start' | 'center' | 'end' | 'stretch'
  justify?: 'start' | 'center' | 'end' | 'between'
  wrap?: boolean
}

const gapMap: Record<GridGap, string> = {
  2: 'gap-2',
  4: 'gap-4',
  6: 'gap-6',
  8: 'gap-8',
  12: 'gap-12',
  16: 'gap-16',
  24: 'gap-24',
}

const alignMap = {
  start: 'items-start',
  center: 'items-center',
  end: 'items-end',
  stretch: 'items-stretch',
} as const

const justifyMap = {
  start: 'justify-start',
  center: 'justify-center',
  end: 'justify-end',
  between: 'justify-between',
} as const

export const Stack = React.forwardRef<HTMLDivElement, StackProps>(
  ({ className, direction = 'vertical', gap = 4, align, justify, wrap, ...props }, ref) => (
    <div
      ref={ref}
      className={cn(
        'flex',
        direction === 'vertical' ? 'flex-col' : 'flex-row',
        gapMap[gap],
        align && alignMap[align],
        justify && justifyMap[justify],
        wrap && 'flex-wrap',
        className,
      )}
      {...props}
    />
  ),
)
Stack.displayName = 'Stack'