Add: wave_limit in combat

tasks/combat/combat.py

@@ -1,3 +1,4 @@
+from module.base.decorator import run_once
 from module.logger import logger
 from tasks.base.assets.assets_base_page import CLOSE
 from tasks.combat.assets.assets_combat_finish import COMBAT_AGAIN, COMBAT_EXIT
@@ -13,18 +14,43 @@ from tasks.map.control.joystick import MapControlJoystick
 class Combat(CombatInteract, CombatPrepare, CombatState, CombatTeam, MapControlJoystick):
     def handle_combat_prepare(self):
         """
+        Returns:
+            bool: If able to run a combat
+
         Pages:
             in: COMBAT_PREPARE
         """
+        self.combat_waves = 1
         current = self.combat_get_trailblaze_power()
         cost = self.combat_get_wave_cost()
         if cost == 10:
-            wave = min(current // self.combat_wave_cost, 6)
-            logger.info(f'Current has {current}, combat costs {self.combat_wave_cost}, able to do {wave} waves')
-            if wave > 0:
-                self.combat_set_wave(wave)
+            # Calyx
+            self.combat_waves = min(current // self.combat_wave_cost, 6)
+            if self.combat_wave_limit:
+                self.combat_waves = min(self.combat_waves, self.combat_wave_limit - self.combat_wave_done)
+                logger.info(
+                    f'Current has {current}, combat costs {self.combat_wave_cost}, '
+                    f'wave={self.combat_wave_done}/{self.combat_wave_limit}, '
+                    f'able to do {self.combat_waves} waves')
             else:
-            logger.info(f'Current has {current}, combat costs {self.combat_wave_cost}, do 1 wave')
+                logger.info(f'Current has {current}, combat costs {self.combat_wave_cost}, '
+                            f'able to do {self.combat_waves} waves')
+            if self.combat_waves > 0:
+                self.combat_set_wave(self.combat_waves)
+        else:
+            # Others
+            logger.info(f'Current has {current}, combat costs {self.combat_wave_cost}, '
+                        f'do {self.combat_waves} wave')
+
+        # Check limits
+        if self.state.TrailblazePower < self.combat_wave_cost:
+            logger.info('Trailblaze power exhausted, cannot continue combat')
+            return False
+        if self.combat_waves <= 0:
+            logger.info('Combat wave limited, cannot continue combat')
+            return False
+
+        return True
 
     def handle_ascension_dungeon_prepare(self):
         """
@@ -72,8 +98,7 @@ class Combat(CombatInteract, CombatPrepare, CombatState, CombatTeam, MapControlJ
                 self.interval_reset(COMBAT_PREPARE)
                 self._map_A_timer.reset()
             if self.appear(COMBAT_PREPARE, interval=2):
-                self.handle_combat_prepare()
-                if self.state.TrailblazePower < self.combat_wave_cost:
+                if not self.handle_combat_prepare():
                     return False
                 self.device.click(COMBAT_PREPARE)
                 self.interval_reset(COMBAT_PREPARE)
@@ -120,8 +145,16 @@ class Combat(CombatInteract, CombatPrepare, CombatState, CombatTeam, MapControlJ
             in: COMBAT_AGAIN
         """
         current = self.combat_get_trailblaze_power(expect_reduce=True)
+        # Wave limit
+        if self.combat_wave_limit:
+            if self.combat_wave_done + self.combat_waves > self.combat_wave_limit:
+                logger.info(f'Combat wave limit: {self.combat_wave_done}/{self.combat_wave_limit}, '
+                            f'can not run again')
+                return False
+
+        # Cost limit
         if self.combat_wave_cost == 10:
-            if current >= self.combat_wave_cost * 6:
+            if current >= self.combat_wave_cost * self.combat_waves:
                 logger.info(f'Current has {current}, combat costs {self.combat_wave_cost}, can run again')
                 return True
             else:
@@ -135,6 +168,24 @@ class Combat(CombatInteract, CombatPrepare, CombatState, CombatTeam, MapControlJ
                 logger.info(f'Current has {current}, combat costs {self.combat_wave_cost}, can not run again')
                 return False
 
+    def _combat_should_reenter(self):
+        """
+        Returns:
+            bool: True to re-enter combat and run with another wave settings
+        """
+        # Wave limit
+        if self.combat_wave_limit:
+            if self.combat_wave_done < self.combat_wave_limit:
+                logger.info(f'Combat wave limit: {self.combat_wave_done}/{self.combat_wave_limit}, '
+                            f'run again with less waves')
+                return True
+            else:
+                return False
+        # Cost limit
+        if self.state.TrailblazePower >= self.combat_wave_cost:
+            logger.info('Still having some trailblaze power run with less waves to empty it')
+            return True
+
     def combat_finish(self) -> bool:
         """
         Returns:
@@ -146,6 +197,12 @@ class Combat(CombatInteract, CombatPrepare, CombatState, CombatTeam, MapControlJ
                 is_combat_executing if again
         """
         logger.hr('Combat finish')
+
+        @run_once
+        def add_wave_done():
+            self.combat_wave_done += self.combat_waves
+            logger.info(f'Done {self.combat_waves} waves at total')
+
         skip_first_screenshot = True
         while 1:
             if skip_first_screenshot:
@@ -164,6 +221,7 @@ class Combat(CombatInteract, CombatPrepare, CombatState, CombatTeam, MapControlJ
             # Click
             # Game client might slow to response COMBAT_AGAIN clicks
             if self.appear(COMBAT_AGAIN, interval=5):
+                add_wave_done()
                 if self._combat_can_again():
                     self.device.click(COMBAT_AGAIN)
                 else:
@@ -202,12 +260,13 @@ class Combat(CombatInteract, CombatPrepare, CombatState, CombatTeam, MapControlJ
                 self.device.click(COMBAT_EXIT)
                 continue
 
-    def combat(self, team: int = 1, skip_first_screenshot=True):
+    def combat(self, team: int = 1, wave_limit: int = 0, skip_first_screenshot=True):
         """
         Combat until trailblaze power runs out.
 
         Args:
             team: 1 to 6.
+            wave_limit: Limit combat runs, 0 means no limit.
             skip_first_screenshot:
 
         Pages:
@@ -218,8 +277,11 @@ class Combat(CombatInteract, CombatPrepare, CombatState, CombatTeam, MapControlJ
         if not skip_first_screenshot:
             self.device.screenshot()
 
+        self.combat_wave_limit = wave_limit
+        self.combat_wave_done = 0
         while 1:
             logger.hr('Combat', level=2)
+            logger.info(f'Combat, team={team}, wave={self.combat_wave_done}/{self.combat_wave_limit}')
             # Prepare
             prepare = self.combat_prepare(team)
             if not prepare:
@@ -229,8 +291,7 @@ class Combat(CombatInteract, CombatPrepare, CombatState, CombatTeam, MapControlJ
             self.combat_execute()
             # Finish
             finish = self.combat_finish()
-            if self.state.TrailblazePower >= self.combat_wave_cost:
-                logger.info('Still having some trailblaze power run with less waves to empty it')
+            if self._combat_should_reenter():
                 continue
             if finish:
                 break

tasks/combat/prepare.py

@@ -23,7 +23,13 @@ class TrailblazePowerOcr(DigitCounter):
 
 
 class CombatPrepare(UI):
-    combat_wave_cost = True
+    # Current combat waves,
+    combat_waves = 1
+    # Limit combat runs, 0 means no limit.
+    combat_wave_limit = 0
+    combat_wave_done = 0
+    # E.g. 10, 30, 40
+    combat_wave_cost = 10
 
     def combat_set_wave(self, count=6):
         """
@@ -82,7 +88,7 @@ class CombatPrepare(UI):
         Get traiblaze power cost and set it to `combat_cost`
 
         Returns:
-            int: 10, 30, 40
+            int: E.g. 10, 30, 40
 
         Pages:
             in: COMBAT_PREPARE