import type { ReactNode, ReactElement } from 'react'
import { memo, Children, cloneElement, isValidElement } from 'react'
import { cn } from '@/lib/utils'

type ImageChildProps = {
  showOverlay?: boolean
  overlayText?: string
}

export interface GeneratedImagesGridProps {
  children: ReactNode
  className?: string
  gridClassName?: string
  isGenerating?: boolean
  maxVisibleWhenGenerating?: number
  formatRemaining?: (remaining: number) => string
}

export const GeneratedImagesGrid = memo(function GeneratedImagesGrid({
  children,
  className = '',
  gridClassName = '',
  isGenerating = false,
  maxVisibleWhenGenerating = 6,
  formatRemaining = (remaining: number) => `+${remaining}`,
}: GeneratedImagesGridProps) {
  const childArray = Children.toArray(children)
  const totalCount = childArray.length

  const shouldTruncate = isGenerating && totalCount > maxVisibleWhenGenerating
  const visibleCount = shouldTruncate ? maxVisibleWhenGenerating : totalCount
  const remainingCount = totalCount - maxVisibleWhenGenerating + 1

  const visibleChildren = childArray.slice(0, visibleCount)

  const processedChildren = visibleChildren.map((child, index) => {
    const isLastAndNeedOverlay = shouldTruncate && index === visibleCount - 1
    if (isLastAndNeedOverlay && isValidElement(child)) {
      return cloneElement(child as ReactElement<ImageChildProps>, {
        showOverlay: true,
        overlayText: formatRemaining(remainingCount),
      })
    }
    return child
  })

  return (
    <div className={cn('mt-5 mb-2', className)}>
      <div className={cn('flex flex-wrap gap-3', gridClassName)}>
        {processedChildren}
      </div>
    </div>
  )
})

GeneratedImagesGrid.displayName = 'GeneratedImagesGrid'