From d01df424488af41238ef540266f32f5f19ce2d22 Mon Sep 17 00:00:00 2001
From: R-J Lim <kgerbil@gmail.com>
Date: Sun, 14 Apr 2024 19:47:39 +0900
Subject: [PATCH] Fix cases where UI hidden for screenshot does not return

---
 .../handlers/video/record-media-handler.ts    | 24 ++++-----
 extension/src/services/binding.ts             |  5 +-
 extension/src/services/image-capturer.ts      | 52 ++++++++++++-------
 3 files changed, 46 insertions(+), 35 deletions(-)

diff --git a/extension/src/handlers/video/record-media-handler.ts b/extension/src/handlers/video/record-media-handler.ts
index f49acf15..de0560b2 100755
--- a/extension/src/handlers/video/record-media-handler.ts
+++ b/extension/src/handlers/video/record-media-handler.ts
@@ -77,6 +77,16 @@ export default class RecordMediaHandler {
                 Math.min(subtitle.end - subtitle.start, recordMediaCommand.message.imageDelay),
                 { maxWidth, maxHeight, rect, frameId }
             );
+            imagePromise.then(() => {
+                const screenshotTakenCommand: ExtensionToVideoCommand<ScreenshotTakenMessage> = {
+                    sender: 'asbplayer-extension-to-video',
+                    message: {
+                        command: 'screenshot-taken',
+                    },
+                    src: recordMediaCommand.src,
+                };
+                chrome.tabs.sendMessage(senderTab.id!, screenshotTakenCommand);
+            });
         }
 
         if (audioPromise) {
@@ -90,22 +100,12 @@ export default class RecordMediaHandler {
             };
         }
 
-        let imageBase64: string | undefined;
-
         if (imagePromise) {
-            imageBase64 = await imagePromise;
-            const screenshotTakenCommand: ExtensionToVideoCommand<ScreenshotTakenMessage> = {
-                sender: 'asbplayer-extension-to-video',
-                message: {
-                    command: 'screenshot-taken',
-                },
-                src: recordMediaCommand.src,
-            };
-            chrome.tabs.sendMessage(senderTab.id!, screenshotTakenCommand);
+            await imagePromise;
 
             // Use the last screenshot taken to allow user to re-take screenshot while audio is recording
             imageModel = {
-                base64: imageBase64!,
+                base64: this._imageCapturer.lastImageBase64!,
                 extension: 'jpeg',
             };
         }
diff --git a/extension/src/services/binding.ts b/extension/src/services/binding.ts
index 2d1664ac..6141e0c5 100755
--- a/extension/src/services/binding.ts
+++ b/extension/src/services/binding.ts
@@ -635,12 +635,11 @@ export default class Binding {
                         const screenshotTakenMessage = request.message as ScreenshotTakenMessage;
                         this.subtitleController.forceHideSubtitles = false;
                         this.mobileVideoOverlayController.forceHide = false;
+                        this.controlsController.show();
 
-                        if (screenshotTakenMessage.ankiUiState) {
+                        if (!this.recordingMedia && screenshotTakenMessage.ankiUiState) {
                             this.ankiUiController.showAfterRetakingScreenshot(this, screenshotTakenMessage.ankiUiState);
                         }
-
-                        this.controlsController.show();
                         break;
                     case 'alert':
                         // ignore
diff --git a/extension/src/services/image-capturer.ts b/extension/src/services/image-capturer.ts
index 594154e5..763676d6 100755
--- a/extension/src/services/image-capturer.ts
+++ b/extension/src/services/image-capturer.ts
@@ -16,6 +16,7 @@ export default class ImageCapturer {
     private imageBase64Reject: ((error: any) => void) | undefined;
     private lastCaptureTimeoutId?: NodeJS.Timeout;
 
+    private _capturing = false;
     private _lastImageBase64?: string;
 
     constructor(settings: SettingsProvider) {
@@ -27,6 +28,10 @@ export default class ImageCapturer {
     }
 
     capture(tabId: number, src: string, delay: number, captureParams: ImageCaptureParams): Promise<string> {
+        if (this._capturing) {
+            return Promise.reject(new Error('Already capturing visible tab'));
+        }
+
         this._lastImageBase64 = undefined;
 
         if (
@@ -56,31 +61,38 @@ export default class ImageCapturer {
         reject: (error: any) => void
     ) {
         const timeoutId = setTimeout(() => {
-            captureVisibleTab(tabId).then(async (dataUrl) => {
-                try {
-                    if (timeoutId !== this.lastCaptureTimeoutId) {
-                        // The promise was already resolved by another call to capture with a shorter delay
-                        return;
+            this._capturing = true;
+            captureVisibleTab(tabId)
+                .then(async (dataUrl) => {
+                    try {
+                        this._capturing = false;
+
+                        if (timeoutId !== this.lastCaptureTimeoutId) {
+                            // The promise was already resolved by another call to capture with a shorter delay
+                            return;
+                        }
+
+                        const croppedDataUrl = await this._cropAndResize(dataUrl, tabId, src, captureParams);
+
+                        if (timeoutId !== this.lastCaptureTimeoutId) {
+                            // The promise was already resolved by another call to capture with a shorter delay
+                            return;
+                        }
+
+                        this._lastImageBase64 = croppedDataUrl.substring(croppedDataUrl.indexOf(',') + 1);
+                        resolve(this._lastImageBase64);
+                    } catch (e) {
+                        reject(e);
                     }
-
-                    const croppedDataUrl = await this._cropAndResize(dataUrl, tabId, src, captureParams);
-
-                    if (timeoutId !== this.lastCaptureTimeoutId) {
-                        // The promise was already resolved by another call to capture with a shorter delay
-                        return;
-                    }
-
-                    this._lastImageBase64 = croppedDataUrl.substring(croppedDataUrl.indexOf(',') + 1);
-                    resolve(this._lastImageBase64);
-                } catch (e) {
-                    reject(e);
-                } finally {
+                })
+                .catch(reject)
+                .finally(() => {
+                    this._capturing = false;
                     this.imageBase64Promise = undefined;
                     this.imageBase64Resolve = undefined;
                     this.imageBase64Reject = undefined;
                     this.lastCaptureTimeoutId = undefined;
-                }
-            });
+                });
         }, delay);
         this.lastCaptureTimeoutId = timeoutId;
     }