From 41cc46ca0ee4994143c659d66e430f5da49e0784 Mon Sep 17 00:00:00 2001
From: Maschell <Maschell@gmx.de>
Date: Thu, 26 Jan 2023 13:23:10 +0100
Subject: [PATCH] Abort taking a screenshot if the target doesn't get rendered

---
 src/function_patcher.cpp | 19 ++++++++++++++++---
 src/retain_vars.cpp      |  4 ++++
 src/retain_vars.hpp      |  4 ++++
 src/utils/utils.cpp      |  9 +++++++++
 4 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/src/function_patcher.cpp b/src/function_patcher.cpp
index 9f909e6..d916990 100644
--- a/src/function_patcher.cpp
+++ b/src/function_patcher.cpp
@@ -46,7 +46,8 @@ void RequestScreenshot() {
         if (gImageSource == IMAGE_SOURCE_TV_AND_DRC || gImageSource == IMAGE_SOURCE_TV) {
             if (gTakeScreenshotTV == SCREENSHOT_STATE_READY) {
                 DEBUG_FUNCTION_LINE("Requested screenshot for TV!");
-                gTakeScreenshotTV = SCREENSHOT_STATE_REQUESTED;
+                gTakeScreenshotTV   = SCREENSHOT_STATE_REQUESTED;
+                gReadySinceFramesTV = 0;
             } else if (!gInProgressNotificationDisplayedTV) {
                 if ((err = NotificationModule_AddErrorNotificationWithCallback("Screenshot of the TV already in progress.",
                                                                                AlreadyInProgressCallback,
@@ -65,7 +66,8 @@ void RequestScreenshot() {
             }
             if (gTakeScreenshotDRC == SCREENSHOT_STATE_READY) {
                 DEBUG_FUNCTION_LINE("Requested screenshot for DRC!");
-                gTakeScreenshotDRC = SCREENSHOT_STATE_REQUESTED;
+                gTakeScreenshotDRC   = SCREENSHOT_STATE_REQUESTED;
+                gReadySinceFramesDRC = 0;
             } else if (!gInProgressNotificationDisplayedDRC) {
                 if ((err = NotificationModule_AddErrorNotificationWithCallback("Screenshot of the GamePad already in progress.",
                                                                                AlreadyInProgressCallback,
@@ -149,7 +151,17 @@ DECL_FUNCTION(void, WPADRead, WPADChan chan, WPADStatusProController *data) {
 
 DECL_FUNCTION(void, GX2CopyColorBufferToScanBuffer, const GX2ColorBuffer *colorBuffer, GX2ScanTarget scan_target) {
     if (gEnabled) {
+        if (gCheckIfScreenRendered) {
+            if (gTakeScreenshotTV == SCREENSHOT_STATE_REQUESTED && ++gReadySinceFramesTV > 5) {
+                gTakeScreenshotTV   = SCREENSHOT_STATE_READY;
+                gReadySinceFramesTV = 0;
+            } else if (gTakeScreenshotDRC == SCREENSHOT_STATE_REQUESTED && ++gReadySinceFramesDRC > 5) {
+                gTakeScreenshotDRC   = SCREENSHOT_STATE_READY;
+                gReadySinceFramesDRC = 0;
+            }
+        }
         if (scan_target == GX2_SCAN_TARGET_TV && colorBuffer != nullptr && gTakeScreenshotTV == SCREENSHOT_STATE_REQUESTED) {
+            gReadySinceFramesTV = 0;
             DEBUG_FUNCTION_LINE("Lets take a screenshot from TV.");
             if (!takeScreenshot((GX2ColorBuffer *) colorBuffer, scan_target, gTVSurfaceFormat, gOutputFormat, gQuality)) {
                 gTakeScreenshotTV = SCREENSHOT_STATE_READY;
@@ -157,6 +169,7 @@ DECL_FUNCTION(void, GX2CopyColorBufferToScanBuffer, const GX2ColorBuffer *colorB
                 gTakeScreenshotTV = SCREENSHOT_STATE_SAVING;
             }
         } else if (scan_target == GX2_SCAN_TARGET_DRC0 && colorBuffer != nullptr && gTakeScreenshotDRC == SCREENSHOT_STATE_REQUESTED) {
+            gReadySinceFramesDRC = 0;
             DEBUG_FUNCTION_LINE("Lets take a screenshot from DRC.");
             if (!takeScreenshot((GX2ColorBuffer *) colorBuffer, scan_target, gDRCSurfaceFormat, gOutputFormat, gQuality)) {
                 gTakeScreenshotDRC = SCREENSHOT_STATE_READY;
@@ -222,4 +235,4 @@ WUPS_MUST_REPLACE(GX2GetCurrentScanBuffer, WUPS_LOADER_LIBRARY_GX2, GX2GetCurren
 WUPS_MUST_REPLACE(GX2CopyColorBufferToScanBuffer, WUPS_LOADER_LIBRARY_GX2, GX2CopyColorBufferToScanBuffer);
 WUPS_MUST_REPLACE(GX2SetTVBuffer, WUPS_LOADER_LIBRARY_GX2, GX2SetTVBuffer);
 WUPS_MUST_REPLACE(GX2SetDRCBuffer, WUPS_LOADER_LIBRARY_GX2, GX2SetDRCBuffer);
-WUPS_MUST_REPLACE(WPADRead, WUPS_LOADER_LIBRARY_PADSCORE, WPADRead);
\ No newline at end of file
+WUPS_MUST_REPLACE(WPADRead, WUPS_LOADER_LIBRARY_PADSCORE, WPADRead);
diff --git a/src/retain_vars.cpp b/src/retain_vars.cpp
index 0ec3e51..398a4e5 100644
--- a/src/retain_vars.cpp
+++ b/src/retain_vars.cpp
@@ -27,3 +27,7 @@ bool gBlockDRCScreenshots = false;
 bool gBlockScreenshots    = false;
 
 bool gInitNotificationModule = false;
+bool gCheckIfScreenRendered  = false;
+
+uint32_t gReadySinceFramesTV  = 0;
+uint32_t gReadySinceFramesDRC = 0;
diff --git a/src/retain_vars.hpp b/src/retain_vars.hpp
index 68ec228..9154ebc 100644
--- a/src/retain_vars.hpp
+++ b/src/retain_vars.hpp
@@ -30,3 +30,7 @@ extern bool gBlockDRCScreenshots;
 extern bool gBlockScreenshots;
 
 extern bool gInitNotificationModule;
+extern bool gCheckIfScreenRendered;
+
+extern uint32_t gReadySinceFramesTV;
+extern uint32_t gReadySinceFramesDRC;
diff --git a/src/utils/utils.cpp b/src/utils/utils.cpp
index c53ebcc..dc8d9b9 100644
--- a/src/utils/utils.cpp
+++ b/src/utils/utils.cpp
@@ -139,4 +139,13 @@ void ApplyGameSpecificPatches() {
     } else {
         gBlockScreenshots = false;
     }
+
+    if (titleID == 0x0005000010138300L || // Donkey Kong Country Tropical Freeze EUR
+        titleID == 0x0005000010144800L || // Donkey Kong Country Tropical Freeze JPN
+        titleID == 0x0005000010137F00L    // Donkey Kong Country Tropical Freeze USA
+    ) {
+        gCheckIfScreenRendered = true;
+    } else {
+        gCheckIfScreenRendered = false;
+    }
 }