import * as React from 'react'
import { Slot } from '@radix-ui/react-slot'
import { cn } from './cva'

type SpaceToken = '1' | '2' | '3' | '4' | '5' | '6' | '8' | '10' | '12'
type InlineAt = 'never' | 'md' | 'lg'
type Align = 'stretch' | 'start' | 'center' | 'end'
type Justify = 'start' | 'center' | 'end' | 'between'
type CSSVars = React.CSSProperties & Record<`--${string}`, string>

export interface StackProps extends Omit<React.HTMLAttributes<HTMLDivElement>, 'style'> {
  asChild?: boolean
  gap?: SpaceToken
  inlineAt?: InlineAt
  align?: Align
  justify?: Justify
  style?: React.CSSProperties
}

const alignClass: Record<Align, string> = {
  stretch: 'items-stretch',
  start: 'items-start',
  center: 'items-center',
  end: 'items-end',
}

const justifyClass: Record<Justify, string> = {
  start: 'justify-start',
  center: 'justify-center',
  end: 'justify-end',
  between: 'justify-between',
}

const inlineClass: Record<InlineAt, string> = {
  never: '',
  md: '@md:flex-row',
  lg: '@lg:flex-row',
}

export const Stack = React.forwardRef<HTMLDivElement, StackProps>(
  ({ asChild, className, gap = '4', inlineAt = 'never', align = 'stretch', justify = 'start', style, ...props }, ref) => {
    const Comp = asChild ? Slot : 'div'
    const stackStyle: CSSVars = {
      '--stack-gap': `var(--mod-space-${gap})`,
      gap: 'var(--stack-gap)',
      minWidth: 0,
      ...style,
    }

    return (
      <Comp
        ref={ref}
        className={cn('@container flex flex-col', alignClass[align], justifyClass[justify], inlineClass[inlineAt], className)}
        style={stackStyle}
        {...props}
      />
    )
  },
)

Stack.displayName = 'Stack'