import { clsx } from 'clsx';
import React from 'react';

type GridCols = 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12;
type GridRows = 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12;
type GridGap = 'none' | 'xs' | 'sm' | 'md' | 'lg' | 'xl' | '2xl' | '3xl';
type GridAlign = 'start' | 'center' | 'end' | 'stretch';
type GridJustify = 'start' | 'center' | 'end' | 'stretch';

interface GridProps extends React.HTMLAttributes<HTMLDivElement> {
  cols?: GridCols | { sm?: GridCols; md?: GridCols; lg?: GridCols; xl?: GridCols };
  rows?: GridRows | { sm?: GridRows; md?: GridRows; lg?: GridRows; xl?: GridRows };
  gap?: GridGap;
  align?: GridAlign;
  justify?: GridJustify;
  autoFlow?: 'row' | 'col' | 'row-dense' | 'col-dense';
  as?: React.ElementType;
  children: React.ReactNode;
}

const GAP_MAP: Record<GridGap, string> = {
  none: 'gap-0',
  xs: 'gap-2',
  sm: 'gap-3',
  md: 'gap-4',
  lg: 'gap-6',
  xl: 'gap-8',
  '2xl': 'gap-12',
  '3xl': 'gap-16',
};

const ALIGN_MAP: Record<GridAlign, string> = {
  start: 'items-start',
  center: 'items-center',
  end: 'items-end',
  stretch: 'items-stretch',
};

const JUSTIFY_MAP: Record<GridJustify, string> = {
  start: 'justify-items-start',
  center: 'justify-items-center',
  end: 'justify-items-end',
  stretch: 'justify-items-stretch',
};

const AUTOFLOW_MAP = {
  row: 'grid-flow-row',
  col: 'grid-flow-col',
  'row-dense': 'grid-flow-row-dense',
  'col-dense': 'grid-flow-col-dense',
};

const getResponsiveClasses = (
  prop:
    | GridCols
    | GridRows
    | {
        sm?: GridCols | GridRows;
        md?: GridCols | GridRows;
        lg?: GridCols | GridRows;
        xl?: GridCols | GridRows;
      }
    | undefined,
  prefix: 'grid-cols' | 'grid-rows'
): string => {
  if (!prop) return '';

  if (typeof prop === 'number') {
    return `${prefix}-${prop}`;
  }

  const classes: string[] = [];
  if (prop.sm) classes.push(`sm:${prefix}-${prop.sm}`);
  if (prop.md) classes.push(`md:${prefix}-${prop.md}`);
  if (prop.lg) classes.push(`lg:${prefix}-${prop.lg}`);
  if (prop.xl) classes.push(`xl:${prefix}-${prop.xl}`);

  return classes.join(' ');
};

const Grid = React.forwardRef<HTMLDivElement, GridProps>(
  (
    {
      cols,
      rows,
      gap = 'md',
      align = 'stretch',
      justify = 'stretch',
      autoFlow,
      as: Component = 'div',
      className,
      children,
      ...props
    },
    ref
  ) => {
    const classes = clsx(
      'grid',
      getResponsiveClasses(cols, 'grid-cols'),
      getResponsiveClasses(rows, 'grid-rows'),
      GAP_MAP[gap],
      ALIGN_MAP[align],
      JUSTIFY_MAP[justify],
      autoFlow && AUTOFLOW_MAP[autoFlow],
      className
    );

    return (
      <Component ref={ref} className={classes} {...props}>
        {children}
      </Component>
    );
  }
);

Grid.displayName = 'Grid';

export default Grid;