import { useCallback, useEffect, useRef, useState } from 'react';

import { Box, CircularProgress, Paper, Stack, Typography } from '@mui/material';
import { type Meta, type StoryObj } from '@storybook/react-vite';

import Table from '../components/design-system/Table';
import TableBody from '../components/design-system/Table/components/TableBody';
import TableCell from '../components/design-system/Table/components/TableCell';
import TableContainer from '../components/design-system/Table/components/TableContainer';
import TableHead from '../components/design-system/Table/components/TableHead';
import TableRow from '../components/design-system/Table/components/TableRow';

import useVirtualizer from './useVirtualizer';

type User = {
  id: number;
  name: string;
  email: string;
  active: string;
};

const USERS: User[] = Array.from({ length: 50 }, (_, index) => ({
  id: index + 1,
  name: `User ${index + 1}`,
  email: `user${index + 1}@example.com`,
  active: index % 6 === 0 ? 'Active' : 'Inactive',
}));

const ROW_HEIGHT_ITEM = 57;
const ROW_HEIGHT_TABLE = 55.4;

// Mock function to simulate fetching more users
const mockFetchUsers = async (
  currentCount: number,
  pageSize: number = 10,
): Promise<User[]> => {
  // Simulate network delay
  await new Promise(resolve => setTimeout(resolve, 1000));

  const newUsers: User[] = Array.from({ length: pageSize }, (_, index) => ({
    id: currentCount + index,
    name: `User ${currentCount + index + 1}`,
    email: `user${currentCount + index + 1}@example.com`,
    active: (currentCount + index) % 6 === 0 ? 'Active' : 'Inactive',
  }));

  return newUsers;
};

type UserItemProps = {
  user: User;
  virtualItem: { start: number };
};

const UserItem = ({ user, virtualItem }: UserItemProps) => (
  <Stack
    id={`${user.id}`}
    sx={{
      alignItems: 'center',
      flexDirection: 'row',
      justifyContent: 'space-between',
      left: 0,
      position: 'absolute',
      top: 0,
      transform: `translateY(${virtualItem.start}px)`,
      width: '100%',
      py: 2,
      borderBottom: '1px solid black',
    }}
  >
    <Typography>{user.name}</Typography>
    <Typography>{user.email}</Typography>
    <Typography sx={{ color: user.active === 'Active' ? 'green' : 'red' }}>
      {user.active}
    </Typography>
  </Stack>
);

const VirtualizedTable = () => {
  const scrollElementRef = useRef<HTMLDivElement>(null);
  const { virtualRows, rowVirtualizer } = useVirtualizer({
    scrollElementRef,
    registers: USERS,
    virtualizerOptions: {
      estimateSize: () => ROW_HEIGHT_TABLE,
      count: USERS.length,
      getScrollElement: () => scrollElementRef.current,
    },
  });

  return (
    <>
      <Typography sx={{ mb: 2 }}>
        Note: MUI Table fixed header breaks after scrolling a bit with Tanstack
        virtualization
      </Typography>
      <TableContainer
        component={Paper}
        containerRef={scrollElementRef}
        sx={{
          height: `${rowVirtualizer.getTotalSize()}px`,
          maxHeight: 300,
          overflow: 'auto',
          width: '100%',
        }}
      >
        <Table stickyHeader>
          <TableHead>
            <TableRow>
              <TableCell headerCell>Name</TableCell>
              <TableCell headerCell>Email</TableCell>
              <TableCell headerCell>Status</TableCell>
            </TableRow>
          </TableHead>
          <TableBody>
            {virtualRows.map((virtualItem, index) => {
              const user = USERS[virtualItem.index];
              return (
                <TableRow
                  id={`${user.id}`}
                  key={user.id}
                  sx={{
                    height: `${virtualItem.size}px`,
                    transform: `translateY(${
                      virtualItem.start - index * virtualItem.size
                    }px)`,
                    backgroundColor: ({ palette }) =>
                      user.active === 'Active'
                        ? palette.new.background.feedback.success
                        : palette.new.background.feedback.error,
                  }}
                >
                  <TableCell>{user.name}</TableCell>
                  <TableCell>{user.email}</TableCell>
                  <TableCell>{user.active}</TableCell>
                </TableRow>
              );
            })}
          </TableBody>
        </Table>
      </TableContainer>
    </>
  );
};

const VirtualizedList = ({ infiniteScroll = false }) => {
  const scrollElementRef = useRef<HTMLElement | null>(null);
  const [users, setUsers] = useState<User[]>(
    infiniteScroll ? USERS.slice(0, 20) : USERS,
  );
  const [isFetchingNextPage, setIsFetchingNextPage] = useState(false);
  const [hasNextPage, setHasNextPage] = useState(true);
  const PAGE_SIZE = 10;
  const MAX_USERS = 100;

  const fetchNextPage = useCallback(async () => {
    if (isFetchingNextPage || !hasNextPage) return;

    setIsFetchingNextPage(true);

    try {
      const newUsers = await mockFetchUsers(users.length, PAGE_SIZE);
      const updatedUsers = [...users, ...newUsers];

      setUsers(updatedUsers);

      // Prevent fetching when we reach the max or no more users
      if (updatedUsers.length >= MAX_USERS) {
        setHasNextPage(false);
      }
    } catch (error) {
      console.error('Error fetching users:', error);
    } finally {
      setIsFetchingNextPage(false);
    }
  }, [users, isFetchingNextPage, hasNextPage]);

  const { virtualRows, rowVirtualizer } = useVirtualizer({
    scrollElementRef,
    registers: users,
    virtualizerOptions: {
      estimateSize: () => ROW_HEIGHT_ITEM,
      count: users.length,
      getScrollElement: () => scrollElementRef.current,
      overscan: 3,
    },
    ...(infiniteScroll && {
      hasNextPage,
      isFetchingNextPage,
      fetchNextPage,
    }),
  });

  return (
    <Box sx={{ width: '100%' }}>
      {infiniteScroll && (
        <Typography sx={{ mb: 2 }}>
          {`Scroll down to load more users. Showing ${users.length}/${MAX_USERS} users`}
        </Typography>
      )}
      <Box
        ref={scrollElementRef}
        sx={{
          height: 300,
          width: '100%',
          overflow: 'auto',
        }}
      >
        <Stack
          sx={{
            height: `${rowVirtualizer.getTotalSize()}px`,
            width: '100%',
            position: 'relative',
          }}
        >
          {virtualRows.map(virtualItem => {
            const user = users[virtualItem.index];
            return (
              <UserItem
                key={user.email}
                user={user}
                virtualItem={virtualItem}
              />
            );
          })}
        </Stack>
        {infiniteScroll && isFetchingNextPage && (
          <Stack
            sx={{
              justifyContent: 'center',
              alignItems: 'center',
              py: 2,
            }}
          >
            <CircularProgress size={20} />
          </Stack>
        )}
        {infiniteScroll && !hasNextPage && users.length > 0 && (
          <Typography
            sx={{
              textAlign: 'center',
              py: 2,
            }}
          >
            No more users to load
          </Typography>
        )}
      </Box>
    </Box>
  );
};

const VirtualizedWindowList = () => {
  const scrollElementRef = useRef<HTMLElement | null>(null);
  const [isScrollElementReady, setIsScrollElementReady] = useState(false);

  // Workaround to get the scroll element ready before the hook is used on the storybook
  // This should be handled by the props "withWindowScroll" and "windowScrollElementId" of useVirtualizer hook
  useEffect(() => {
    const scrollElement = document.getElementById('custom-scroll-container');
    if (scrollElement) {
      scrollElementRef.current = scrollElement;
      setIsScrollElementReady(true);
    }
  }, []);

  const { virtualRows, rowVirtualizer } = useVirtualizer({
    scrollElementRef,
    registers: USERS,
    virtualizerOptions: {
      estimateSize: () => ROW_HEIGHT_ITEM,
      count: USERS.length,
      getScrollElement: () => scrollElementRef.current,
      overscan: 3,
    },
  });

  if (!isScrollElementReady) {
    return null;
  }

  return (
    <Box sx={{ width: '60%' }}>
      <Stack
        sx={{
          height: `${rowVirtualizer.getTotalSize()}px`,
          width: '100%',
          position: 'relative',
        }}
      >
        {virtualRows.map(virtualItem => {
          const user = USERS[virtualItem.index];
          return (
            <UserItem
              key={user.email}
              user={user}
              virtualItem={virtualItem}
            />
          );
        })}
      </Stack>
    </Box>
  );
};

const meta: Meta = {
  title: 'Hooks/useVirtualizer',
  tags: ['autodocs'],
  parameters: {
    docs: {
      description: {
        component:
          'A hook that provides row virtualization on given element. It returns virtualized rows and the virtualizer instance.',
      },
    },
  },
};

export default meta;

type Story = StoryObj<typeof meta>;

export const Default: Story = {
  render: () => <VirtualizedList />,
  parameters: {
    docs: {
      description: {
        story: 'Virtualized list with fixed element quantity',
      },
    },
  },
};

export const WithInfiniteScroll: Story = {
  render: () => <VirtualizedList infiniteScroll />,
  parameters: {
    docs: {
      description: {
        story: 'Virtualized list with infinite scroll',
      },
    },
  },
};

export const WithWindowScroll: Story = {
  render: () => <VirtualizedWindowList />,
  decorators: [
    Story => {
      return (
        <Stack
          id="custom-scroll-container"
          sx={{
            alignItems: 'center',
            height: '400px',
            overflow: 'auto',
            position: 'relative',
            width: '100%',
          }}
        >
          <Typography
            sx={{
              backgroundColor: '#F5F6F8',
              position: 'sticky',
              top: 0,
              zIndex: 1,
            }}
          >
            Scroll down over the list padding on the sides to test the outer
            container scrolling
          </Typography>
          <Story />
        </Stack>
      );
    },
  ],
  parameters: {
    docs: {
      description: {
        story: 'Virtualized list with window / custom parent container scroll',
      },
    },
  },
};

export const TableElement: Story = {
  render: () => <VirtualizedTable />,
  parameters: {
    docs: {
      description: {
        story: 'Virtualized table with fixed element quantity',
      },
    },
  },
};