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

/**
 * WARNING: DO NOT VIOLATE THE PURPOSE OF THIS PRIMITIVE.
 * 
 * Stack is for flexbox-based layouts (stacking elements).
 * 
 * - DO NOT add positioning props (absolute, top, zIndex).
 * - DO NOT add grid props.
 * - DO NOT add background/border props unless it's a specific styled stack.
 * 
 * If you need a more specific layout, create a new component in apps/website/components.
 */

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

export interface StackProps<T extends ElementType> extends Omit<BoxProps<T>, 'children'> {
  as?: T;
  children?: ReactNode;
  className?: string;
  direction?: 'row' | 'col' | { base?: 'row' | 'col'; md?: 'row' | 'col'; lg?: 'row' | 'col' };
  gap?: number | string | ResponsiveGap;
  align?: 'start' | 'center' | 'end' | 'stretch' | 'baseline' | ResponsiveValue<'start' | 'center' | 'end' | 'stretch' | 'baseline'>;
  justify?: 'start' | 'center' | 'end' | 'between' | 'around' | ResponsiveValue<'start' | 'center' | 'end' | 'between' | 'around'>;
  wrap?: boolean;
}

export const Stack = forwardRef(<T extends ElementType = 'div'>(
  { 
    children, 
    className = '', 
    direction = 'col',
    gap = 4,
    align,
    justify,
    wrap = false,
    as,
    ...props
  }: StackProps<T>,
  ref: ForwardedRef<HTMLElement>
) => {
  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 | string | ResponsiveGap | undefined) => {
    if (value === undefined) return '';
    if (typeof value === 'object') {
      const classes = [];
      if (value.base !== undefined) classes.push(typeof value.base === 'number' ? gapClasses[value.base] : `gap-${value.base}`);
      if (value.sm !== undefined) classes.push(typeof value.sm === 'number' ? `sm:${gapClasses[value.sm]}` : `sm:gap-${value.sm}`);
      if (value.md !== undefined) classes.push(typeof value.md === 'number' ? `md:${gapClasses[value.md]}` : `md:gap-${value.md}`);
      if (value.lg !== undefined) classes.push(typeof value.lg === 'number' ? `lg:${gapClasses[value.lg]}` : `lg:gap-${value.lg}`);
      if (value.xl !== undefined) classes.push(typeof value.xl === 'number' ? `xl:${gapClasses[value.xl]}` : `xl:gap-${value.xl}`);
      return classes.join(' ');
    }
    if (typeof value === 'number') return gapClasses[value];
    return `gap-${value}`;
  };

  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',
    wrap ? 'flex-wrap' : '',
    className
  ].filter(Boolean).join(' ');

  const getAlignItemsClass = (value: StackProps<ElementType>['align']) => {
    if (!value) return '';
    const map: Record<string, string> = { start: 'items-start', center: 'items-center', end: 'items-end', stretch: 'items-stretch', baseline: 'items-baseline' };
    if (typeof value === 'object') {
      const classes = [];
      if (value.base) classes.push(map[value.base]);
      if (value.sm) classes.push(`sm:${map[value.sm]}`);
      if (value.md) classes.push(`md:${map[value.md]}`);
      if (value.lg) classes.push(`lg:${map[value.lg]}`);
      if (value.xl) classes.push(`xl:${map[value.xl]}`);
      if (value['2xl']) classes.push(`2xl:${map[value['2xl']]}`);
      return classes.join(' ');
    }
    return map[value];
  };

  const getJustifyContentClass = (value: StackProps<ElementType>['justify']) => {
    if (!value) return '';
    const map: Record<string, string> = { start: 'justify-start', center: 'justify-center', end: 'justify-end', between: 'justify-between', around: 'justify-around' };
    if (typeof value === 'object') {
      const classes = [];
      if (value.base) classes.push(map[value.base]);
      if (value.sm) classes.push(`sm:${map[value.sm]}`);
      if (value.md) classes.push(`md:${map[value.md]}`);
      if (value.lg) classes.push(`lg:${map[value.lg]}`);
      if (value.xl) classes.push(`xl:${map[value.xl]}`);
      if (value['2xl']) classes.push(`2xl:${map[value['2xl']]}`);
      return classes.join(' ');
    }
    return map[value];
  };

  const layoutClasses = [
    getAlignItemsClass(align),
    getJustifyContentClass(justify)
  ].filter(Boolean).join(' ');

  return (
    <Box 
      as={as} 
      ref={ref}
      className={`${classes} ${layoutClasses}`}
      {...props}
    >
      {children}
    </Box>
  );
});

Stack.displayName = 'Stack';