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

export interface ImageGridImage {
  src: string;
  alt?: string;
}

export interface ImageGridProps extends React.HTMLAttributes<HTMLDivElement> {
  images: ImageGridImage[];
  columns?: 1 | 2 | 3 | 4;
  onImageClick?: (index: number) => void;
  loading?: boolean;
}

const columnClasses = {
  1: 'grid-cols-1',
  2: 'grid-cols-2',
  3: 'grid-cols-3',
  4: 'grid-cols-4',
} as const;

const ImageGrid = React.forwardRef<HTMLDivElement, ImageGridProps>(
  ({ images, columns = 2, onImageClick, loading, className, ...props }, ref) => {
    const [loadedImages, setLoadedImages] = React.useState<Set<number>>(new Set());

    const handleImageLoad = React.useCallback((index: number) => {
      setLoadedImages((prev) => new Set(prev).add(index));
    }, []);

    if (loading) {
      return (
        <div
          ref={ref}
          className={cn('grid gap-2', columnClasses[columns], className)}
          {...props}
        >
          {Array.from({ length: columns * 2 }).map((_, i) => (
            <div
              key={i}
              className="aspect-square rounded-lg bg-[var(--surface-200)] animate-pulse"
            />
          ))}
        </div>
      );
    }

    if (images.length === 0) {
      return null;
    }

    return (
      <div
        ref={ref}
        className={cn('grid gap-2', columnClasses[columns], className)}
        {...props}
      >
        {images.map((image, index) => (
          <button
            key={index}
            type="button"
            onClick={() => onImageClick?.(index)}
            disabled={!onImageClick}
            className={cn(
              'relative aspect-square rounded-lg overflow-hidden bg-[var(--surface-200)]',
              'focus:outline-none focus:ring-2 focus:ring-[var(--accent)] focus:ring-offset-2 focus:ring-offset-[var(--background)]',
              onImageClick && 'cursor-pointer hover:opacity-90 transition-opacity'
            )}
          >
            {!loadedImages.has(index) && (
              <div className="absolute inset-0 animate-pulse bg-[var(--surface-300)]" />
            )}
            <img
              src={image.src}
              alt={image.alt || `Generated image ${index + 1}`}
              className={cn(
                'w-full h-full object-cover transition-opacity duration-300',
                loadedImages.has(index) ? 'opacity-100' : 'opacity-0'
              )}
              onLoad={() => handleImageLoad(index)}
            />
          </button>
        ))}
      </div>
    );
  }
);
ImageGrid.displayName = 'ImageGrid';

export { ImageGrid };
