import React, { useRef } from 'react'
import { useVirtualizer, type Virtualizer } from '@tanstack/react-virtual'

type VirtualProps = {
  /** This is the number of rows in the dataset */
  count: number
  /** Children is a function that has access to the useVirtualizer hook value and the ref that needs to be wrapped around the virtualized content. */
  children: ({
    virtualizer,
    virtualizedParentRef,
  }: {
    virtualizer: ReturnType<typeof useVirtualizer>
    virtualizedParentRef: React.RefObject<HTMLDivElement | null>
  }) => React.JSX.Element
  /** Optionally pass in more react-virtual props for the useVirtualizer */
  virtualizerProps?: Parameters<typeof useVirtualizer>[0]
  /** Optionally define the row height for the rows. Note: This is not on a row-by-row basis. It will apply to every row. The default is 40px. */
  rowHeight?: number
}

const Virtual = ({
  count,
  children,
  virtualizerProps,
  rowHeight = 40,
}: VirtualProps): React.JSX.Element => {
  const virtualizedParentRef = useRef(null)
  const virtualizer = useVirtualizer({
    count,
    getScrollElement: () => virtualizedParentRef.current,
    estimateSize: () => rowHeight,
    overscan: 5,
    ...virtualizerProps,
  })

  return children({ virtualizer, virtualizedParentRef })
}

type VirtualListProps = {
  virtualizer: ReturnType<typeof useVirtualizer>
  children: ({
    index,
    measureElement,
  }: {
    index: number
    measureElement: Virtualizer<Element, Element>['measureElement']
  }) => React.JSX.Element
}

export const VirtualList = ({
  virtualizer,
  children,
}: VirtualListProps): React.JSX.Element => {
  return (
    <div
      style={{
        height: `${virtualizer.getTotalSize()}px`,
        width: '100%',
        position: 'relative',
      }}
    >
      {virtualizer.getVirtualItems().map((virtualRow) => (
        <div
          key={virtualRow.index}
          style={{
            transform: `translateY(${virtualRow.start}px)`,
            position: 'absolute',
            top: 0,
            left: 0,
            width: '100%',
          }}
        >
          {children({
            index: virtualRow.index,
            measureElement: virtualizer.measureElement,
          })}
        </div>
      ))}
    </div>
  )
}

Virtual.List = VirtualList

export { Virtual }