From 2e07945d9b73001b1013416abb300be8a0915ab7 Mon Sep 17 00:00:00 2001
From: ArthurKun <16458204+ArthurKun21@users.noreply.github.com>
Date: Wed, 2 Oct 2024 11:14:51 +0800
Subject: [PATCH] wait for AP regen on the last run

---
 .../scripts/entrypoints/AutoBattle.kt         |  7 +-
 .../scripts/modules/Refill.kt                 | 66 +++++++++++++------
 2 files changed, 53 insertions(+), 20 deletions(-)

diff --git a/scripts/src/main/java/io/github/fate_grand_automata/scripts/entrypoints/AutoBattle.kt b/scripts/src/main/java/io/github/fate_grand_automata/scripts/entrypoints/AutoBattle.kt
index b9b88266a..3f3fbd505 100644
--- a/scripts/src/main/java/io/github/fate_grand_automata/scripts/entrypoints/AutoBattle.kt
+++ b/scripts/src/main/java/io/github/fate_grand_automata/scripts/entrypoints/AutoBattle.kt
@@ -531,6 +531,11 @@ class AutoBattle @Inject constructor(
             inventoryFull -> throw BattleExitException(ExitReason.InventoryFull)
         }
 
-        refill.refill()
+        val isLastRun = prefs.selectedServerConfigPref.shouldLimitRuns &&
+                (prefs.selectedServerConfigPref.limitRuns - state.runs) <= 1
+
+        refill.refill(
+            isLastRun = isLastRun
+        )
     }
 }
diff --git a/scripts/src/main/java/io/github/fate_grand_automata/scripts/modules/Refill.kt b/scripts/src/main/java/io/github/fate_grand_automata/scripts/modules/Refill.kt
index 05218f317..8dcca9c77 100644
--- a/scripts/src/main/java/io/github/fate_grand_automata/scripts/modules/Refill.kt
+++ b/scripts/src/main/java/io/github/fate_grand_automata/scripts/modules/Refill.kt
@@ -4,6 +4,7 @@ import io.github.fate_grand_automata.scripts.IFgoAutomataApi
 import io.github.fate_grand_automata.scripts.Images
 import io.github.fate_grand_automata.scripts.ScriptNotify
 import io.github.fate_grand_automata.scripts.entrypoints.AutoBattle
+import io.github.fate_grand_automata.scripts.prefs.IPerServerConfigPrefs
 import io.github.lib_automata.dagger.ScriptScope
 import javax.inject.Inject
 import kotlin.time.Duration.Companion.seconds
@@ -19,34 +20,61 @@ class Refill @Inject constructor(
      * Refills the AP with apples depending on preferences.
      * If needed, loops and wait for AP regeneration
      */
-    private fun refillOnce() {
+    private fun refillOnce(
+        isLastRun: Boolean = false
+    ) {
         val perServerConfigPref = prefs.selectedServerConfigPref
 
-        if (perServerConfigPref.resources.isNotEmpty()
-            && timesRefilled < perServerConfigPref.currentAppleCount
-        ) {
-            //TODO check for OK image between each resource
-            perServerConfigPref.resources
-                .flatMap { locations.locate(it) }
-                .forEach { it.click() }
+        when {
+            /**
+             * If the user has resources to refill and has the wait for AP regen option enabled
+             * and this is the last run, wait for AP regen instead of refilling.
+             */
+            perServerConfigPref.waitForAPRegen && isLastRun -> waitForAPRegen()
+            /**
+             * If the user has resources to refill and has not reached the current apple count,
+             */
+            perServerConfigPref.resources.isNotEmpty() && timesRefilled < perServerConfigPref.currentAppleCount -> {
+                refillAP(perServerConfigPref = perServerConfigPref)
+            }
+            /**
+             * wait for AP regen if the user has the wait for AP regen option enabled
+             */
 
-            1.seconds.wait()
-            locations.staminaOkClick.click()
-            ++timesRefilled
+            perServerConfigPref.waitForAPRegen -> waitForAPRegen()
 
-            3.seconds.wait()
-        } else if (perServerConfigPref.waitForAPRegen) {
-            locations.staminaCloseClick.click()
+            else -> throw AutoBattle.BattleExitException(AutoBattle.ExitReason.APRanOut)
+        }
+    }
+
+    private fun refillAP(perServerConfigPref: IPerServerConfigPrefs) {
+        //TODO check for OK image between each resource
+        perServerConfigPref.resources
+            .flatMap { locations.locate(it) }
+            .forEach { it.click() }
+
+        1.seconds.wait()
+        locations.staminaOkClick.click()
+        ++timesRefilled
+
+        3.seconds.wait()
+    }
+
+    private fun waitForAPRegen() {
+        locations.staminaCloseClick.click()
 
-            messages.notify(ScriptNotify.WaitForAPRegen())
+        messages.notify(ScriptNotify.WaitForAPRegen())
 
-            60.seconds.wait()
-        } else throw AutoBattle.BattleExitException(AutoBattle.ExitReason.APRanOut)
+        60.seconds.wait()
     }
 
-    fun refill() {
+    fun refill(
+        isLastRun: Boolean = false
+    ) {
         if (images[Images.Stamina] in locations.staminaScreenRegion) {
-            refillOnce()
+            refillOnce(
+                isLastRun = isLastRun
+            )
         }
     }