import {
  Children,
  forwardRef,
  type HTMLAttributes,
  type ReactElement,
  useCallback,
  useRef,
  useState,
} from 'react';
import { List, type RowComponentProps } from 'react-window';

const DEFAULT_ROW_HEIGHT = 48;
const MAX_HEIGHT = 240;
const OVERSCAN_COUNT = 5;

type VirtualRowProps = {
  items: ReactElement[];
  onRowMeasured: (index: number, height: number) => void;
};

const VirtualRow = ({
  index,
  style,
  items,
  onRowMeasured,
}: RowComponentProps<VirtualRowProps>) => {
  const item = items[index];
  const measureRef = useCallback(
    (node: HTMLDivElement | null) => {
      if (node) {
        const measured = node.getBoundingClientRect().height;
        if (measured > 0) {
          onRowMeasured(index, measured);
        }
      }
    },
    [index, onRowMeasured],
  );

  return (
    <div style={style}>
      <div ref={measureRef}>{item}</div>
    </div>
  );
};

const VirtualizedListbox = forwardRef<
  HTMLDivElement,
  HTMLAttributes<HTMLElement>
>(({ children, ...props }, ref) => {
  const items = Children.toArray(children) as ReactElement[];
  const rowHeightsRef = useRef<Record<number, number>>({});
  const [measureGeneration, setMeasureGeneration] = useState(0);

  // biome-ignore lint/correctness/useExhaustiveDependencies: measureGeneration creates a new function identity to force react-window to recompute row positions
  const getRowHeight = useCallback(
    (index: number) => rowHeightsRef.current[index] ?? DEFAULT_ROW_HEIGHT,
    [measureGeneration],
  );

  const onRowMeasured = useCallback((index: number, height: number) => {
    if (rowHeightsRef.current[index] !== height) {
      rowHeightsRef.current[index] = height;
      setMeasureGeneration(g => g + 1);
    }
  }, []);

  return (
    <div
      ref={ref}
      {...props}
      style={{ ...props.style, overflow: 'hidden' }}
    >
      <List
        rowCount={items.length}
        rowHeight={getRowHeight}
        style={{ maxHeight: MAX_HEIGHT }}
        overscanCount={OVERSCAN_COUNT}
        rowComponent={VirtualRow}
        rowProps={{ items, onRowMeasured }}
      />
    </div>
  );
});

VirtualizedListbox.displayName = 'VirtualizedListbox';

export default VirtualizedListbox;
