import React, { useState, useCallback, memo, useMemo } from 'react';
import {
  View,
  StyleSheet,
  LayoutChangeEvent,
} from 'react-native';
import { FlashList } from '@shopify/flash-list';

export type MasonryListProps<T> = {
  data: T[];
  renderItem: (item: T, index: number) => React.ReactNode;
  keyExtractor: (item: T) => string;
  getItemHeight?: (item: T, width: number) => number | undefined;
  numColumns?: number;
  gap?: number;
  estimatedItemHeight?: number;
  onEndReached?: () => void;
  onEndReachedThreshold?: number;
  onRefresh?: () => void;
  refreshing?: boolean;
  contentContainerStyle?: object;
};

type RowItem<T> = {
  type: 'row';
  items: T[];
  rowIndex: number;
};

function MasonryListInner<T>({
  data,
  renderItem,
  keyExtractor,
  getItemHeight = () => undefined,
  numColumns = 2,
  gap = 8,
  estimatedItemHeight = 200,
  onEndReached,
  onEndReachedThreshold = 0.5,
  onRefresh,
  refreshing = false,
  contentContainerStyle,
}: MasonryListProps<T>) {
  const [containerWidth, setContainerWidth] = useState(0);

  const handleLayout = useCallback((event: LayoutChangeEvent) => {
    const { width } = event.nativeEvent.layout;
    setContainerWidth(width);
  }, []);

  const rows = useMemo(() => {
    const result: RowItem<T>[] = [];
    for (let i = 0; i < data.length; i += numColumns) {
      result.push({
        type: 'row',
        items: data.slice(i, i + numColumns),
        rowIndex: i / numColumns,
      });
    }
    return result;
  }, [data, numColumns]);

  const itemWidth = useMemo(() => {
    if (containerWidth === 0) return 0;
    return (containerWidth - gap * (numColumns - 1)) / numColumns;
  }, [containerWidth, gap, numColumns]);

  const renderRow = useCallback(({ item: row }: { item: RowItem<T> }) => {
    return (
      <View style={[styles.row, { marginBottom: gap }]}>
        {row.items.map((item, index) => {
          const height = getItemHeight(item, itemWidth) || estimatedItemHeight;
          return (
            <View
              key={keyExtractor(item)}
              style={[
                styles.item,
                {
                  width: itemWidth,
                  height,
                  marginRight: index < row.items.length - 1 ? gap : 0,
                  backgroundColor: '#2E3031',
                  borderRadius: 12,
                  overflow: 'hidden',
                },
              ]}
            >
              {renderItem(item, row.rowIndex * numColumns + index)}
            </View>
          );
        })}
      </View>
    );
  }, [renderItem, keyExtractor, itemWidth, gap, numColumns, getItemHeight, estimatedItemHeight]);

  if (containerWidth === 0) {
    return <View style={styles.container} onLayout={handleLayout} />;
  }

  return (
    <View style={styles.container} onLayout={handleLayout}>
      <FlashList
        data={rows}
        renderItem={renderRow}
        keyExtractor={(item) => `row-${item.rowIndex}`}
        onEndReached={onEndReached}
        onEndReachedThreshold={onEndReachedThreshold}
        onRefresh={onRefresh}
        refreshing={refreshing}
        contentContainerStyle={contentContainerStyle}
      />
    </View>
  );
}

export const MasonryList = memo(MasonryListInner) as typeof MasonryListInner;

const styles = StyleSheet.create({
  container: {
    flex: 1,
  },
  row: {
    flexDirection: 'row',
  },
  item: {},
});