From d8d61970087764bf027f95557fc110c55c356f3e Mon Sep 17 00:00:00 2001
From: Reece <reece@stirlingpdf.com>
Date: Fri, 26 Sep 2025 17:48:41 +0100
Subject: [PATCH] fix page count issue

---
 frontend/src/tools/Sign.tsx               |  5 +++--
 frontend/src/utils/signatureFlattening.ts | 22 +++++++++-------------
 2 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/frontend/src/tools/Sign.tsx b/frontend/src/tools/Sign.tsx
index 0706cc6b5..0304928a0 100644
--- a/frontend/src/tools/Sign.tsx
+++ b/frontend/src/tools/Sign.tsx
@@ -17,7 +17,7 @@ const Sign = (props: BaseToolProps) => {
   const { setWorkbench } = useNavigation();
   const { setSignatureConfig, activateDrawMode, activateSignaturePlacementMode, deactivateDrawMode, updateDrawSettings, undo, redo, signatureApiRef, getImageData, setSignaturesApplied } = useSignature();
   const { consumeFiles, selectors } = useFileContext();
-  const { exportActions } = useViewer();
+  const { exportActions, getScrollState } = useViewer();
 
   // Track which signature mode was active for reactivation after save
   const activeModeRef = useRef<'draw' | 'placement' | null>(null);
@@ -76,7 +76,8 @@ const Sign = (props: BaseToolProps) => {
         exportActions,
         selectors,
         consumeFiles,
-        originalFile
+        originalFile,
+        getScrollState
       });
 
       if (success) {
diff --git a/frontend/src/utils/signatureFlattening.ts b/frontend/src/utils/signatureFlattening.ts
index ffd86b2e3..44310376f 100644
--- a/frontend/src/utils/signatureFlattening.ts
+++ b/frontend/src/utils/signatureFlattening.ts
@@ -19,10 +19,11 @@ interface SignatureFlatteningOptions {
   selectors: MinimalFileContextSelectors;
   consumeFiles: (inputFileIds: FileId[], outputStirlingFiles: StirlingFile[], outputStirlingFileStubs: StirlingFileStub[]) => Promise<FileId[]>;
   originalFile?: StirlingFile;
+  getScrollState: () => { currentPage: number; totalPages: number };
 }
 
 export async function flattenSignatures(options: SignatureFlatteningOptions): Promise<boolean> {
-  const { signatureApiRef, getImageData, exportActions, selectors, consumeFiles, originalFile } = options;
+  const { signatureApiRef, getImageData, exportActions, selectors, consumeFiles, originalFile, getScrollState } = options;
 
   try {
     // Step 1: Extract all annotations from EmbedPDF before export
@@ -31,12 +32,13 @@ export async function flattenSignatures(options: SignatureFlatteningOptions): Pr
     if (signatureApiRef?.current) {
       console.log('Extracting annotations from all pages...');
 
-      // Dynamically check all pages until we encounter consecutive errors
-      let pageIndex = 0;
-      let consecutiveErrors = 0;
-      const maxConsecutiveErrors = 3; // Stop after 3 consecutive page access failures
+      // Get actual page count from viewer
+      const scrollState = getScrollState();
+      const totalPages = scrollState.totalPages;
+      console.log(`Document has ${totalPages} pages`);
 
-      while (consecutiveErrors < maxConsecutiveErrors) {
+      // Check only actual pages that exist in the document
+      for (let pageIndex = 0; pageIndex < totalPages; pageIndex++) {
         try {
           const pageAnnotations = await signatureApiRef.current.getPageAnnotations(pageIndex);
           if (pageAnnotations && pageAnnotations.length > 0) {
@@ -66,15 +68,9 @@ export async function flattenSignatures(options: SignatureFlatteningOptions): Pr
               console.log(`Found ${sessionAnnotations.length} session annotations on page ${pageIndex + 1} (out of ${pageAnnotations.length} total)`);
             }
           }
-
-          // Reset consecutive error count on successful page access
-          consecutiveErrors = 0;
         } catch (pageError) {
-          consecutiveErrors++;
-          console.warn(`Error extracting annotations from page ${pageIndex + 1} (error ${consecutiveErrors}/${maxConsecutiveErrors}):`, pageError);
+          console.warn(`Error extracting annotations from page ${pageIndex + 1}:`, pageError);
         }
-
-        pageIndex++;
       }
     }