import { memo, useMemo } from 'react';
import { View, StyleSheet, type ViewStyle } from 'react-native';

type WaterfallItem = {
  w: number;
  h: number;
  data: any;
};

type WaterfallProps<T> = {
  items: WaterfallItem[];
  column?: number;
  itemPadding?: number;
  renderItem: (item: T, width: number, height: number) => React.ReactNode;
  containerWidth: number;
};

function WaterfallComponent<T>({
  items,
  column = 2,
  itemPadding = 8,
  renderItem,
  containerWidth,
}: WaterfallProps<T>) {
  const layout = useMemo(() => {
    const columns: { h: number; items: Array<WaterfallItem & { x: number; y: number }> }[] =
      Array.from({ length: column }, () => ({ h: 0, items: [] }));

    const itemWidth = (containerWidth - itemPadding * (column - 1)) / column;

    items.forEach((item) => {
      const minColumn = columns.reduce((min, col, idx) =>
        col.h < columns[min].h ? idx : min, 0);

      const col = columns[minColumn];
      const layoutItem = {
        ...item,
        x: minColumn * (itemWidth + itemPadding),
        y: col.h > 0 ? col.h + itemPadding : 0,
        w: itemWidth,
      };

      col.h = layoutItem.y + item.h + 40;
      col.items.push(layoutItem);
    });

    const maxHeight = Math.max(...columns.map(col => col.h));
    const allItems = columns.flatMap(col => col.items);

    return { items: allItems, height: maxHeight };
  }, [items, column, itemPadding, containerWidth]);

  return (
    <View style={[styles.container, { height: layout.height }]}>
      {layout.items.map((item, index) => (
        <View
          key={item.data.id || index}
          style={[
            styles.item,
            {
              position: 'absolute',
              left: item.x,
              top: item.y,
              width: item.w,
              height: item.h + 48
            },
          ]}
        >
          {renderItem(item.data, item.w, item.h)}
        </View>
      ))}
    </View>
  );
}

const styles = StyleSheet.create({
  container: {
    position: 'relative',
    width: '100%',
  },
  item: {
    overflow: 'hidden',
  },
});

export const Waterfall = memo(WaterfallComponent) as typeof WaterfallComponent;