From fd0cfe97bde524b794d2b1b3ff560889a92ff5ce Mon Sep 17 00:00:00 2001
From: Tanish Manku <tanish2k09.dev@gmail.com>
Date: Mon, 23 Mar 2026 20:52:51 -0700
Subject: [PATCH] add arrow key navigation to image generation lightbox
 (#gen-view-arrow-key)

- Left/Right navigates carousel slides (instant jump, loops)
- Up/Down navigates between generation batches
- Opening a specific image correctly highlights it in the carousel via initialIndex
- Add galleryModalLightboxNavPrev/Next signals so TextToImage can provide its
  own bundle navigation callbacks independently of the gallery's item list
- Fix: gallery nav callbacks were undefined for TextToImage-opened lightboxes
  since generated batches are not in the gallery's flatFilteredItems
---
 .../app/src/pages/PageImage/TextToImage.tsx   | 100 +++++++++++-------
 .../gallery-modal/src/lib/gallery-modal.tsx   |   9 +-
 .../src/lib/galleryModalSignals.ts            |   5 +
 .../lightbox-modal/src/lib/lightbox-modal.tsx |  38 ++++---
 4 files changed, 95 insertions(+), 57 deletions(-)

diff --git a/frontend/apps/artcraft/app/src/pages/PageImage/TextToImage.tsx b/frontend/apps/artcraft/app/src/pages/PageImage/TextToImage.tsx
index 8a66448847..0b89ef9aec 100644
--- a/frontend/apps/artcraft/app/src/pages/PageImage/TextToImage.tsx
+++ b/frontend/apps/artcraft/app/src/pages/PageImage/TextToImage.tsx
@@ -1,4 +1,4 @@
-import { useEffect, useMemo, useRef, useState } from "react";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import {
   JobContextType,
   getThumbnailUrl,
@@ -31,6 +31,8 @@ import {
   galleryModalLightboxImage,
   galleryModalLightboxMediaId,
   galleryModalLightboxVisible,
+  galleryModalLightboxNavPrev,
+  galleryModalLightboxNavNext,
 } from "@storyteller/ui-gallery-modal";
 import { Badge } from "@storyteller/ui-badge";
 import { twMerge } from "tailwind-merge";
@@ -125,6 +127,63 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => {
   // Like Midjourney instead of a "chat history" style.
   const inverseBatch = [...batches].reverse();
 
+  // Keep a ref so nav callbacks always see the latest batch list without
+  // needing to be recreated on every render.
+  const inverseBatchRef = useRef(inverseBatch);
+  inverseBatchRef.current = inverseBatch;
+
+  const openBatchInLightbox = useCallback(
+    (targetBatch: (typeof inverseBatch)[0], startIndex = 0) => {
+      const images = targetBatch.images;
+      if (!images.length) return;
+
+      const imageUrls = images.map(
+        (img) =>
+          getThumbnailUrl(img.maybe_thumbnail_template, { width: 3200 }) ??
+          img.cdn_url,
+      );
+      const mediaTokens = images.map((img) => img.media_token);
+      const actionUrls = images.map((img) => img.cdn_url);
+
+      const clickedImg = images[startIndex];
+      const batchIndex = inverseBatchRef.current.findIndex(
+        (b) => b.id === targetBatch.id,
+      );
+
+      galleryModalLightboxNavPrev.value =
+        batchIndex > 0
+          ? () => openBatchInLightbox(inverseBatchRef.current[batchIndex - 1])
+          : null;
+      galleryModalLightboxNavNext.value =
+        batchIndex < inverseBatchRef.current.length - 1
+          ? () => openBatchInLightbox(inverseBatchRef.current[batchIndex + 1])
+          : null;
+
+      galleryModalLightboxMediaId.value = clickedImg.media_token;
+      galleryModalLightboxImage.value = {
+        id: clickedImg.media_token,
+        label: targetBatch.prompt || "Generated Image",
+        thumbnail:
+          getThumbnailUrl(clickedImg.maybe_thumbnail_template, {
+            width: THUMBNAIL_SIZES.MEDIUM,
+          }) ?? clickedImg.cdn_url,
+        fullImage:
+          getThumbnailUrl(clickedImg.maybe_thumbnail_template, {
+            width: 2048,
+          }) ?? clickedImg.cdn_url,
+        createdAt: new Date(targetBatch.createdAt).toISOString(),
+        mediaClass: "image" as const,
+        mediaTokens,
+        imageUrls,
+        actionUrls,
+        initialIndex: startIndex,
+        thumbnailUrlTemplate: clickedImg.maybe_thumbnail_template,
+      };
+      galleryModalLightboxVisible.value = true;
+    },
+    [],
+  );
+
   return (
     <div
       ref={containerRef}
@@ -223,43 +282,8 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => {
                             <button
                               key={img.media_token}
                               onClick={() => {
-                                const lightboxItem = {
-                                  id: img.media_token,
-                                  label: batch.prompt || "Generated Image",
-                                  thumbnail:
-                                    getThumbnailUrl(
-                                      img.maybe_thumbnail_template,
-                                      { width: THUMBNAIL_SIZES.MEDIUM },
-                                    ) ?? img.cdn_url,
-                                  fullImage:
-                                    getThumbnailUrl(
-                                      img.maybe_thumbnail_template,
-                                      { width: 2048 },
-                                    ) ?? img.cdn_url,
-                                  createdAt: new Date(
-                                    batch.createdAt,
-                                  ).toISOString(),
-                                  mediaClass: "image" as const,
-                                  mediaTokens: batch.images.map(
-                                    (image) => image.media_token,
-                                  ),
-                                  imageUrls: batch.images.map(
-                                    (image) =>
-                                      getThumbnailUrl(
-                                        image.maybe_thumbnail_template,
-                                        { width: 3200 },
-                                      ) ?? image.cdn_url,
-                                  ),
-                                  actionUrls: batch.images.map(
-                                    (image) => image.cdn_url,
-                                  ),
-                                  thumbnailUrlTemplate:
-                                    img.maybe_thumbnail_template,
-                                };
-                                galleryModalLightboxMediaId.value =
-                                  lightboxItem.id;
-                                galleryModalLightboxImage.value = lightboxItem;
-                                galleryModalLightboxVisible.value = true;
+                                const startIndex = batch.images.indexOf(img);
+                                openBatchInLightbox(batch, startIndex);
                               }}
                               className="aspect-square w-full overflow-hidden rounded-lg transition-opacity duration-200 hover:cursor-pointer hover:opacity-75"
                             >
diff --git a/frontend/libs/components/gallery-modal/src/lib/gallery-modal.tsx b/frontend/libs/components/gallery-modal/src/lib/gallery-modal.tsx
index f94734d7af..5ec7af23b0 100644
--- a/frontend/libs/components/gallery-modal/src/lib/gallery-modal.tsx
+++ b/frontend/libs/components/gallery-modal/src/lib/gallery-modal.tsx
@@ -32,6 +32,8 @@ import {
   galleryModalLightboxMediaId,
   galleryModalLightboxImage,
   galleryModalLightboxVisible,
+  galleryModalLightboxNavPrev,
+  galleryModalLightboxNavNext,
 } from "./galleryModalSignals";
 import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import {
@@ -687,6 +689,8 @@ export const GalleryModal = React.memo(
       lightboxVisibleSignal.value = false;
       lightboxImageSignal.value = null;
       galleryModalLightboxMediaId.value = null;
+      galleryModalLightboxNavPrev.value = null;
+      galleryModalLightboxNavNext.value = null;
     }, []);
 
     const handleCloseGallery = useCallback(() => {
@@ -1279,6 +1283,7 @@ export const GalleryModal = React.memo(
             imageUrls={imageUrls}
             actionUrls={actionUrls}
             mediaTokens={(lightboxImageSignal.value as any)?.mediaTokens}
+            initialIndex={(lightboxImageSignal.value as any)?.initialIndex}
             batchImageToken={
               (lightboxImageSignal.value as any)?.batchImageToken
             }
@@ -1296,8 +1301,8 @@ export const GalleryModal = React.memo(
             onRemoveBackgroundClicked={onRemoveBackgroundClicked}
             onMake3DObjectClicked={onMake3DObjectClicked}
             onMake3DWorldClicked={onMake3DWorldClicked}
-            onNavigatePrev={handleNavigatePrev}
-            onNavigateNext={handleNavigateNext}
+            onNavigatePrev={galleryModalLightboxNavPrev.value ?? handleNavigatePrev}
+            onNavigateNext={galleryModalLightboxNavNext.value ?? handleNavigateNext}
             onNavigateToMedia={handleNavigateToMedia}
           />
         )}
diff --git a/frontend/libs/components/gallery-modal/src/lib/galleryModalSignals.ts b/frontend/libs/components/gallery-modal/src/lib/galleryModalSignals.ts
index 6fd8bcd032..c14a05b7d2 100644
--- a/frontend/libs/components/gallery-modal/src/lib/galleryModalSignals.ts
+++ b/frontend/libs/components/gallery-modal/src/lib/galleryModalSignals.ts
@@ -9,3 +9,8 @@ export const galleryModalVisibleViewMode = signal(false);
 export const galleryModalLightboxMediaId = signal<string | null>(null);
 export const galleryModalLightboxVisible = signal(false);
 export const galleryModalLightboxImage = signal<any>(null);
+// Custom nav callbacks – set by pages that manage their own item list (e.g.
+// TextToImage). Null when opened from the gallery browse, in which case the
+// gallery's own computed handlers are used instead.
+export const galleryModalLightboxNavPrev = signal<(() => void) | null>(null);
+export const galleryModalLightboxNavNext = signal<(() => void) | null>(null);
diff --git a/frontend/libs/components/lightbox-modal/src/lib/lightbox-modal.tsx b/frontend/libs/components/lightbox-modal/src/lib/lightbox-modal.tsx
index a09abff988..135ec947e0 100644
--- a/frontend/libs/components/lightbox-modal/src/lib/lightbox-modal.tsx
+++ b/frontend/libs/components/lightbox-modal/src/lib/lightbox-modal.tsx
@@ -94,6 +94,7 @@ interface LightboxModalProps {
   onNavigatePrev?: () => void;
   onNavigateNext?: () => void;
   onNavigateToMedia?: (mediaToken: string) => void;
+  initialIndex?: number;
 }
 
 export function LightboxModal({
@@ -123,6 +124,7 @@ export function LightboxModal({
   onNavigatePrev,
   onNavigateNext,
   onNavigateToMedia,
+  initialIndex,
 }: LightboxModalProps) {
   // NB(bt,2025-06-14): We add ?cors=1 to the image url to prevent caching "sec-fetch-mode: no-cors" from
   // the <image> tag request from being cached. If we then drag it into the canvas after it's been cached,
@@ -408,9 +410,10 @@ export function LightboxModal({
   }, [emblaMainApi, onSelect]);
 
   useEffect(() => {
-    setSelectedIndex(0);
-    emblaMainApi?.scrollTo(0, true);
-    emblaThumbsApi?.scrollTo(0, true);
+    const idx = initialIndex ?? 0;
+    setSelectedIndex(idx);
+    emblaMainApi?.scrollTo(idx, true);
+    emblaThumbsApi?.scrollTo(idx, true);
   }, [batchImageToken, imageUrl, emblaMainApi, emblaThumbsApi]);
 
   const selectedImageUrl = effectiveImageUrls[selectedIndex] ?? null;
@@ -456,33 +459,32 @@ export function LightboxModal({
     [mediaClass, batchImages],
   );
 
-  // Keyboard navigation for carousel slides and gallery prev/next
+  // Keyboard navigation:
+  //   Left / Right  →  carousel slide (loop wraps automatically)
+  //   Up / Down     →  previous / next bundle
   useEffect(() => {
     if (!isOpen) return;
     const handleKeyDown = (e: KeyboardEvent) => {
-      // Don't navigate if user is interacting with an input, textarea, etc.
       const tag = (e.target as HTMLElement)?.tagName;
       if (tag === "INPUT" || tag === "TEXTAREA" || tag === "SELECT") return;
 
       if (e.key === "ArrowLeft") {
         e.preventDefault();
-        if (emblaMainApi && selectedIndex > 0) {
-          emblaMainApi.scrollPrev(true); // jump = instant for keyboard nav
-        } else if (onNavigatePrev) {
-          onNavigatePrev();
-        }
+        emblaMainApi?.scrollPrev(true);
       } else if (e.key === "ArrowRight") {
         e.preventDefault();
-        if (emblaMainApi && selectedIndex < effectiveImageUrls.length - 1) {
-          emblaMainApi.scrollNext(true); // jump = instant for keyboard nav
-        } else if (onNavigateNext) {
-          onNavigateNext();
-        }
+        emblaMainApi?.scrollNext(true);
+      } else if (e.key === "ArrowUp") {
+        e.preventDefault();
+        onNavigatePrev?.();
+      } else if (e.key === "ArrowDown") {
+        e.preventDefault();
+        onNavigateNext?.();
       }
     };
     window.addEventListener("keydown", handleKeyDown);
     return () => window.removeEventListener("keydown", handleKeyDown);
-  }, [isOpen, emblaMainApi, selectedIndex, effectiveImageUrls, onNavigatePrev, onNavigateNext]);
+  }, [isOpen, emblaMainApi, onNavigatePrev, onNavigateNext]);
 
   return (
     <>
@@ -504,7 +506,9 @@ export function LightboxModal({
         </Modal.DragHandle>
 
         {/* content grid */}
-        <div className="flex h-full gap-4">
+        <div
+          className="flex h-full gap-4"
+        >
           {/* image panel - flexible width */}
           <div className="group/nav relative flex h-full flex-1 items-center justify-center overflow-hidden rounded-l-xl bg-black/30">
             {!selectedImageUrl ? (