From dd0b4dc744813a8615419744e0f75657712e1624 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Fri, 21 Jul 2023 15:37:20 +1000
Subject: [PATCH] fix(ui): fix next prev buttons

---
 .../components/ImageGrid/GalleryImageGrid.tsx | 18 +++++---
 .../gallery/hooks/useNextPrevImage.ts         |  2 +-
 .../src/services/api/hooks/useBoardTotal.ts   | 44 +++++--------------
 3 files changed, 24 insertions(+), 40 deletions(-)

diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImageGrid.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImageGrid.tsx
index a43f9ce07b..e4b996fc96 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImageGrid.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImageGrid.tsx
@@ -19,6 +19,7 @@ import GalleryImage from './GalleryImage';
 import ImageGridItemContainer from './ImageGridItemContainer';
 import ImageGridListContainer from './ImageGridListContainer';
 import { selectListImagesBaseQueryArgs } from 'features/gallery/store/gallerySelectors';
+import { useBoardTotal } from 'services/api/hooks/useBoardTotal';
 
 const overlayScrollbarsConfig: UseOverlayScrollbarsParams = {
   defer: true,
@@ -40,7 +41,10 @@ const GalleryImageGrid = () => {
   const [initialize, osInstance] = useOverlayScrollbars(
     overlayScrollbarsConfig
   );
-
+  const selectedBoardId = useAppSelector(
+    (state) => state.gallery.selectedBoardId
+  );
+  const { currentViewTotal } = useBoardTotal(selectedBoardId);
   const queryArgs = useAppSelector(selectListImagesBaseQueryArgs);
 
   const { currentData, isFetching, isSuccess, isError } =
@@ -49,19 +53,23 @@ const GalleryImageGrid = () => {
   const [listImages] = useLazyListImagesQuery();
 
   const areMoreAvailable = useMemo(() => {
-    if (!currentData) {
+    if (!currentData || !currentViewTotal) {
       return false;
     }
-    return currentData.ids.length < currentData.total;
-  }, [currentData]);
+    return currentData.ids.length < currentViewTotal;
+  }, [currentData, currentViewTotal]);
 
   const handleLoadMoreImages = useCallback(() => {
+    if (!areMoreAvailable) {
+      return;
+    }
+
     listImages({
       ...queryArgs,
       offset: currentData?.ids.length ?? 0,
       limit: IMAGE_LIMIT,
     });
-  }, [listImages, queryArgs, currentData?.ids.length]);
+  }, [areMoreAvailable, listImages, queryArgs, currentData?.ids.length]);
 
   useEffect(() => {
     // Initialize the gallery's custom scrollbar
diff --git a/invokeai/frontend/web/src/features/gallery/hooks/useNextPrevImage.ts b/invokeai/frontend/web/src/features/gallery/hooks/useNextPrevImage.ts
index 57db29c6f1..1dc0e2c10d 100644
--- a/invokeai/frontend/web/src/features/gallery/hooks/useNextPrevImage.ts
+++ b/invokeai/frontend/web/src/features/gallery/hooks/useNextPrevImage.ts
@@ -64,7 +64,7 @@ export const nextPrevImageButtonsSelector = createSelector(
       isOnFirstImage: currentImageIndex === 0,
       isOnLastImage:
         !isNaN(currentImageIndex) && currentImageIndex === imagesLength - 1,
-      areMoreImagesAvailable: data?.total ?? 0 > imagesLength,
+      areMoreImagesAvailable: (data?.total ?? 0) > imagesLength,
       isFetching: status === 'pending',
       nextImage,
       prevImage,
diff --git a/invokeai/frontend/web/src/services/api/hooks/useBoardTotal.ts b/invokeai/frontend/web/src/services/api/hooks/useBoardTotal.ts
index e693ff87d0..1a9e69ff2d 100644
--- a/invokeai/frontend/web/src/services/api/hooks/useBoardTotal.ts
+++ b/invokeai/frontend/web/src/services/api/hooks/useBoardTotal.ts
@@ -1,45 +1,21 @@
+import { useAppSelector } from 'app/store/storeHooks';
+import { BoardId } from 'features/gallery/store/gallerySlice';
+import { useMemo } from 'react';
 import {
-  BoardId,
-  INITIAL_IMAGE_LIMIT,
-} from 'features/gallery/store/gallerySlice';
-import {
-  ListImagesArgs,
   useGetBoardAssetsTotalQuery,
   useGetBoardImagesTotalQuery,
 } from '../endpoints/images';
 
-const baseQueryArgs: ListImagesArgs = {
-  offset: 0,
-  limit: INITIAL_IMAGE_LIMIT,
-  is_intermediate: false,
-};
-
 export const useBoardTotal = (board_id: BoardId) => {
+  const galleryView = useAppSelector((state) => state.gallery.galleryView);
+
   const { data: totalImages } = useGetBoardImagesTotalQuery(board_id);
   const { data: totalAssets } = useGetBoardAssetsTotalQuery(board_id);
-  // const imagesQueryArg = useMemo(() => {
-  //   const categories = IMAGE_CATEGORIES;
-  //   return { board_id, categories, ...baseQueryArgs };
-  // }, [board_id]);
 
-  // const assetsQueryArg = useMemo(() => {
-  //   const categories = ASSETS_CATEGORIES;
-  //   return { board_id, categories, ...baseQueryArgs };
-  // }, [board_id]);
+  const currentViewTotal = useMemo(
+    () => (galleryView === 'images' ? totalImages : totalAssets),
+    [galleryView, totalAssets, totalImages]
+  );
 
-  // const { total: totalImages } = useListImagesQuery(
-  //   imagesQueryArg ?? skipToken,
-  //   {
-  //     selectFromResult: ({ currentData }) => ({ total: currentData?.total }),
-  //   }
-  // );
-
-  // const { total: totalAssets } = useListImagesQuery(
-  //   assetsQueryArg ?? skipToken,
-  //   {
-  //     selectFromResult: ({ currentData }) => ({ total: currentData?.total }),
-  //   }
-  // );
-
-  return { totalImages, totalAssets };
+  return { totalImages, totalAssets, currentViewTotal };
 };