import React, { ReactNode, ElementType } from 'react';
import { Box, BoxProps } from './Box';

type Spacing = 0 | 0.5 | 1 | 1.5 | 2 | 2.5 | 3 | 3.5 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 14 | 16 | 20 | 24 | 28 | 32 | 36 | 40 | 44 | 48 | 52 | 56 | 60 | 64 | 72 | 80 | 96;

interface ResponsiveGap {
  base?: number;
  sm?: number;
  md?: number;
  lg?: number;
  xl?: number;
}

interface StackProps<T extends ElementType> extends Omit<BoxProps<T>, 'children' | 'className' | 'gap'> {
  children: ReactNode;
  className?: string;
  direction?: 'row' | 'col' | { base?: 'row' | 'col'; md?: 'row' | 'col'; lg?: 'row' | 'col' };
  gap?: number | ResponsiveGap;
  align?: 'start' | 'center' | 'end' | 'stretch' | 'baseline';
  justify?: 'start' | 'center' | 'end' | 'between' | 'around';
  wrap?: boolean;
  center?: boolean;
  m?: Spacing | any;
  mt?: Spacing | any;
  mb?: Spacing | any;
  ml?: Spacing | any;
  mr?: Spacing | any;
  p?: Spacing | any;
  pt?: Spacing | any;
  pb?: Spacing | any;
  pl?: Spacing | any;
  pr?: Spacing | any;
  px?: Spacing | any;
  py?: Spacing | any;
  rounded?: 'none' | 'sm' | 'md' | 'lg' | 'xl' | '2xl' | 'full';
  style?: React.CSSProperties;
  role?: string;
  'aria-label'?: string;
  'aria-live'?: 'polite' | 'assertive' | 'off';
}

export function Stack<T extends ElementType = 'div'>({ 
  children, 
  className = '', 
  direction = 'col',
  gap = 4,
  align = 'stretch',
  justify = 'start',
  wrap = false,
  center = false,
  m, mt, mb, ml, mr,
  p, pt, pb, pl, pr, px, py,
  rounded,
  as,
  ...props
}: StackProps<T>) {
  const gapClasses: Record<number, string> = {
    0: 'gap-0',
    1: 'gap-1',
    2: 'gap-2',
    3: 'gap-3',
    4: 'gap-4',
    5: 'gap-5',
    6: 'gap-6',
    8: 'gap-8',
    10: 'gap-10',
    12: 'gap-12',
    16: 'gap-16'
  };

  const getGapClasses = (value: number | ResponsiveGap | undefined) => {
    if (value === undefined) return '';
    if (typeof value === 'object') {
      const classes = [];
      if (value.base !== undefined) classes.push(gapClasses[value.base]);
      if (value.sm !== undefined) classes.push(`sm:${gapClasses[value.sm]}`);
      if (value.md !== undefined) classes.push(`md:${gapClasses[value.md]}`);
      if (value.lg !== undefined) classes.push(`lg:${gapClasses[value.lg]}`);
      if (value.xl !== undefined) classes.push(`xl:${gapClasses[value.xl]}`);
      return classes.join(' ');
    }
    return gapClasses[value];
  };

  const spacingMap: Record<number, string> = {
    0: '0', 0.5: '0.5', 1: '1', 1.5: '1.5', 2: '2', 2.5: '2.5', 3: '3', 3.5: '3.5', 4: '4',
    5: '5', 6: '6', 7: '7', 8: '8', 9: '9', 10: '10', 11: '11', 12: '12', 14: '14',
    16: '16', 20: '20', 24: '24', 28: '28', 32: '32', 36: '36', 40: '40', 44: '44',
    48: '48', 52: '52', 56: '56', 60: '60', 64: '64', 72: '72', 80: '80', 96: '96'
  };

  const roundedClasses = {
    none: 'rounded-none',
    sm: 'rounded-sm',
    md: 'rounded-md',
    lg: 'rounded-lg',
    xl: 'rounded-xl',
    '2xl': 'rounded-2xl',
    full: 'rounded-full'
  };

  const classes = [
    'flex',
    typeof direction === 'string' 
      ? (direction === 'col' ? 'flex-col' : 'flex-row')
      : [
          direction.base === 'col' ? 'flex-col' : (direction.base === 'row' ? 'flex-row' : ''),
          direction.md === 'col' ? 'md:flex-col' : (direction.md === 'row' ? 'md:flex-row' : ''),
          direction.lg === 'col' ? 'lg:flex-col' : (direction.lg === 'row' ? 'lg:flex-row' : ''),
        ].filter(Boolean).join(' '),
    getGapClasses(gap) || 'gap-4',
    center ? 'items-center justify-center' : `items-${align} justify-${justify}`,
    wrap ? 'flex-wrap' : '',
    m !== undefined ? `m-${spacingMap[m]}` : '',
    mt !== undefined ? `mt-${spacingMap[mt]}` : '',
    mb !== undefined ? `mb-${spacingMap[mb]}` : '',
    ml !== undefined ? `ml-${spacingMap[ml]}` : '',
    mr !== undefined ? `mr-${spacingMap[mr]}` : '',
    p !== undefined ? `p-${spacingMap[p]}` : '',
    pt !== undefined ? `pt-${spacingMap[pt]}` : '',
    pb !== undefined ? `pb-${spacingMap[pb]}` : '',
    pl !== undefined ? `pl-${spacingMap[pl]}` : '',
    pr !== undefined ? `pr-${spacingMap[pr]}` : '',
    px !== undefined ? `px-${spacingMap[px]}` : '',
    py !== undefined ? `py-${spacingMap[py]}` : '',
    rounded ? roundedClasses[rounded] : '',
    className
  ].filter(Boolean).join(' ');

  return <Box as={as} className={classes} {...props}>{children}</Box>;
}