Refactor: Abstract StaminaStatus class to handle various states

tasks/combat/assets/assets_combat_fuel.py

@@ -1,165 +0,0 @@
-from module.base.button import Button, ButtonWrapper
-
-# This file was auto-generated, do not modify it manually. To generate:
-# ``` python -m dev_tools.button_extract ```
-
-EXTRACT_RESERVED_TRAILBLAZE_POWER = ButtonWrapper(
-    name='EXTRACT_RESERVED_TRAILBLAZE_POWER',
-    share=Button(
-        file='./assets/share/combat/fuel/EXTRACT_RESERVED_TRAILBLAZE_POWER.png',
-        area=(909, 506, 929, 526),
-        search=(889, 486, 949, 546),
-        color=(91, 91, 91),
-        button=(909, 506, 929, 526),
-    ),
-)
-FUEL = ButtonWrapper(
-    name='FUEL',
-    share=Button(
-        file='./assets/share/combat/fuel/FUEL.png',
-        area=(592, 276, 688, 366),
-        search=(474, 271, 811, 396),
-        color=(123, 96, 134),
-        button=(592, 276, 688, 366),
-    ),
-)
-FUEL_ENTRANCE = ButtonWrapper(
-    name='FUEL_ENTRANCE',
-    share=Button(
-        file='./assets/share/combat/fuel/FUEL_ENTRANCE.png',
-        area=(1035, 26, 1056, 48),
-        search=(1015, 6, 1076, 68),
-        color=(188, 180, 226),
-        button=(1035, 26, 1056, 48),
-    ),
-)
-FUEL_MINUS = ButtonWrapper(
-    name='FUEL_MINUS',
-    share=Button(
-        file='./assets/share/combat/fuel/FUEL_MINUS.png',
-        area=(472, 425, 510, 450),
-        search=(452, 405, 530, 470),
-        color=(236, 236, 236),
-        button=(472, 425, 510, 450),
-    ),
-)
-FUEL_PLUS = ButtonWrapper(
-    name='FUEL_PLUS',
-    share=Button(
-        file='./assets/share/combat/fuel/FUEL_PLUS.png',
-        area=(967, 426, 1005, 449),
-        search=(947, 406, 1025, 469),
-        color=(232, 232, 232),
-        button=(967, 426, 1005, 449),
-    ),
-)
-FUEL_SELECTED = ButtonWrapper(
-    name='FUEL_SELECTED',
-    share=Button(
-        file='./assets/share/combat/fuel/FUEL_SELECTED.png',
-        area=(587, 271, 692, 368),
-        search=(474, 271, 811, 396),
-        color=(136, 112, 144),
-        button=(587, 271, 692, 368),
-    ),
-)
-FUEL_SLIDER = ButtonWrapper(
-    name='FUEL_SLIDER',
-    share=Button(
-        file='./assets/share/combat/fuel/FUEL_SLIDER.png',
-        area=(561, 434, 916, 441),
-        search=(541, 414, 936, 461),
-        color=(215, 185, 154),
-        button=(561, 434, 916, 441),
-    ),
-)
-OCR_EXTRACT_RESERVED_TRAILBLAZE_POWER_COUNT = ButtonWrapper(
-    name='OCR_EXTRACT_RESERVED_TRAILBLAZE_POWER_COUNT',
-    share=Button(
-        file='./assets/share/combat/fuel/OCR_EXTRACT_RESERVED_TRAILBLAZE_POWER_COUNT.png',
-        area=(425, 415, 688, 436),
-        search=(405, 395, 708, 456),
-        color=(192, 192, 192),
-        button=(425, 415, 688, 436),
-    ),
-)
-OCR_FUEL = ButtonWrapper(
-    name='OCR_FUEL',
-    share=Button(
-        file='./assets/share/combat/fuel/OCR_FUEL.png',
-        area=(605, 369, 678, 386),
-        search=(585, 349, 698, 406),
-        color=(66, 66, 66),
-        button=(605, 369, 678, 386),
-    ),
-)
-OCR_FUEL_COUNT = ButtonWrapper(
-    name='OCR_FUEL_COUNT',
-    share=Button(
-        file='./assets/share/combat/fuel/OCR_FUEL_COUNT.png',
-        area=(686, 409, 881, 425),
-        search=(666, 389, 901, 445),
-        color=(205, 205, 205),
-        button=(686, 409, 881, 425),
-    ),
-)
-OCR_RESERVED_TRAILBLAZE_POWER = ButtonWrapper(
-    name='OCR_RESERVED_TRAILBLAZE_POWER',
-    share=Button(
-        file='./assets/share/combat/fuel/OCR_RESERVED_TRAILBLAZE_POWER.png',
-        area=(883, 29, 992, 44),
-        search=(863, 9, 1012, 64),
-        color=(51, 65, 65),
-        button=(883, 29, 992, 44),
-    ),
-)
-RESERVED_MINUS = ButtonWrapper(
-    name='RESERVED_MINUS',
-    share=Button(
-        file='./assets/share/combat/fuel/RESERVED_MINUS.png',
-        area=(248, 474, 281, 498),
-        search=(228, 454, 301, 518),
-        color=(238, 238, 238),
-        button=(248, 474, 281, 498),
-    ),
-)
-RESERVED_PLUS = ButtonWrapper(
-    name='RESERVED_PLUS',
-    share=Button(
-        file='./assets/share/combat/fuel/RESERVED_PLUS.png',
-        area=(938, 475, 974, 498),
-        search=(918, 455, 994, 518),
-        color=(232, 232, 232),
-        button=(938, 475, 974, 498),
-    ),
-)
-RESERVED_SLIDER = ButtonWrapper(
-    name='RESERVED_SLIDER',
-    share=Button(
-        file='./assets/share/combat/fuel/RESERVED_SLIDER.png',
-        area=(334, 483, 873, 489),
-        search=(314, 463, 893, 509),
-        color=(212, 173, 130),
-        button=(334, 483, 873, 489),
-    ),
-)
-RESERVED_TRAILBLAZE_POWER_ENTRANCE = ButtonWrapper(
-    name='RESERVED_TRAILBLAZE_POWER_ENTRANCE',
-    share=Button(
-        file='./assets/share/combat/fuel/RESERVED_TRAILBLAZE_POWER_ENTRANCE.png',
-        area=(895, 26, 916, 48),
-        search=(875, 6, 936, 68),
-        color=(154, 213, 214),
-        button=(895, 26, 916, 48),
-    ),
-)
-USING_FUEL = ButtonWrapper(
-    name='USING_FUEL',
-    share=Button(
-        file='./assets/share/combat/fuel/USING_FUEL.png',
-        area=(263, 265, 363, 365),
-        search=(243, 245, 383, 385),
-        color=(161, 116, 129),
-        button=(263, 265, 363, 365),
-    ),
-)

tasks/combat/assets/assets_combat_prepare.py

@@ -20,16 +20,6 @@ COMBAT_PREPARE = ButtonWrapper(
         button=(956, 640, 1225, 676),
     ),
 )
-OCR_TRAILBLAZE_POWER = ButtonWrapper(
-    name='OCR_TRAILBLAZE_POWER',
-    share=Button(
-        file='./assets/share/combat/prepare/OCR_TRAILBLAZE_POWER.png',
-        area=(998, 26, 1130, 48),
-        search=(978, 6, 1150, 68),
-        color=(77, 76, 87),
-        button=(998, 26, 1130, 48),
-    ),
-)
 OCR_WAVE_COST = ButtonWrapper(
     name='OCR_WAVE_COST',
     share=Button(

tasks/combat/assets/assets_combat_stamina_fuel.py

@@ -0,0 +1,85 @@
+from module.base.button import Button, ButtonWrapper
+
+# This file was auto-generated, do not modify it manually. To generate:
+# ``` python -m dev_tools.button_extract ```
+
+FUEL = ButtonWrapper(
+    name='FUEL',
+    share=Button(
+        file='./assets/share/combat/stamina/fuel/FUEL.png',
+        area=(592, 276, 688, 366),
+        search=(474, 271, 811, 396),
+        color=(123, 96, 134),
+        button=(592, 276, 688, 366),
+    ),
+)
+FUEL_MINUS = ButtonWrapper(
+    name='FUEL_MINUS',
+    share=Button(
+        file='./assets/share/combat/stamina/fuel/FUEL_MINUS.png',
+        area=(472, 425, 510, 450),
+        search=(452, 405, 530, 470),
+        color=(236, 236, 236),
+        button=(472, 425, 510, 450),
+    ),
+)
+FUEL_PLUS = ButtonWrapper(
+    name='FUEL_PLUS',
+    share=Button(
+        file='./assets/share/combat/stamina/fuel/FUEL_PLUS.png',
+        area=(967, 426, 1005, 449),
+        search=(947, 406, 1025, 469),
+        color=(232, 232, 232),
+        button=(967, 426, 1005, 449),
+    ),
+)
+FUEL_SELECTED = ButtonWrapper(
+    name='FUEL_SELECTED',
+    share=Button(
+        file='./assets/share/combat/stamina/fuel/FUEL_SELECTED.png',
+        area=(587, 271, 692, 368),
+        search=(474, 271, 811, 396),
+        color=(136, 112, 144),
+        button=(587, 271, 692, 368),
+    ),
+)
+FUEL_SLIDER = ButtonWrapper(
+    name='FUEL_SLIDER',
+    share=Button(
+        file='./assets/share/combat/stamina/fuel/FUEL_SLIDER.png',
+        area=(561, 434, 916, 441),
+        search=(541, 414, 936, 461),
+        color=(215, 185, 154),
+        button=(561, 434, 916, 441),
+    ),
+)
+OCR_FUEL = ButtonWrapper(
+    name='OCR_FUEL',
+    share=Button(
+        file='./assets/share/combat/stamina/fuel/OCR_FUEL.png',
+        area=(605, 369, 678, 386),
+        search=(585, 349, 698, 406),
+        color=(66, 66, 66),
+        button=(605, 369, 678, 386),
+    ),
+)
+OCR_FUEL_COUNT = ButtonWrapper(
+    name='OCR_FUEL_COUNT',
+    share=Button(
+        file='./assets/share/combat/stamina/fuel/OCR_FUEL_COUNT.png',
+        area=(686, 409, 881, 425),
+        search=(666, 389, 901, 445),
+        color=(205, 205, 205),
+        button=(686, 409, 881, 425),
+    ),
+)
+USING_FUEL = ButtonWrapper(
+    name='USING_FUEL',
+    share=Button(
+        file='./assets/share/combat/stamina/fuel/USING_FUEL.png',
+        area=(263, 265, 363, 365),
+        search=(243, 245, 383, 385),
+        color=(161, 116, 129),
+        button=(263, 265, 363, 365),
+    ),
+)

tasks/combat/assets/assets_combat_stamina_reserved.py

@@ -0,0 +1,55 @@
+from module.base.button import Button, ButtonWrapper
+
+# This file was auto-generated, do not modify it manually. To generate:
+# ``` python -m dev_tools.button_extract ```
+
+EXTRACT_RESERVED_TRAILBLAZE_POWER = ButtonWrapper(
+    name='EXTRACT_RESERVED_TRAILBLAZE_POWER',
+    share=Button(
+        file='./assets/share/combat/stamina/reserved/EXTRACT_RESERVED_TRAILBLAZE_POWER.png',
+        area=(909, 506, 929, 526),
+        search=(889, 486, 949, 546),
+        color=(91, 91, 91),
+        button=(909, 506, 929, 526),
+    ),
+)
+OCR_EXTRACT_RESERVED_TRAILBLAZE_POWER_COUNT = ButtonWrapper(
+    name='OCR_EXTRACT_RESERVED_TRAILBLAZE_POWER_COUNT',
+    share=Button(
+        file='./assets/share/combat/stamina/reserved/OCR_EXTRACT_RESERVED_TRAILBLAZE_POWER_COUNT.png',
+        area=(425, 415, 688, 436),
+        search=(405, 395, 708, 456),
+        color=(192, 192, 192),
+        button=(425, 415, 688, 436),
+    ),
+)
+RESERVED_MINUS = ButtonWrapper(
+    name='RESERVED_MINUS',
+    share=Button(
+        file='./assets/share/combat/stamina/reserved/RESERVED_MINUS.png',
+        area=(248, 474, 281, 498),
+        search=(228, 454, 301, 518),
+        color=(238, 238, 238),
+        button=(248, 474, 281, 498),
+    ),
+)
+RESERVED_PLUS = ButtonWrapper(
+    name='RESERVED_PLUS',
+    share=Button(
+        file='./assets/share/combat/stamina/reserved/RESERVED_PLUS.png',
+        area=(938, 475, 974, 498),
+        search=(918, 455, 994, 518),
+        color=(232, 232, 232),
+        button=(938, 475, 974, 498),
+    ),
+)
+RESERVED_SLIDER = ButtonWrapper(
+    name='RESERVED_SLIDER',
+    share=Button(
+        file='./assets/share/combat/stamina/reserved/RESERVED_SLIDER.png',
+        area=(334, 483, 873, 489),
+        search=(314, 463, 893, 509),
+        color=(212, 173, 130),
+        button=(334, 483, 873, 489),
+    ),
+)

tasks/combat/assets/assets_combat_stamina_status.py

@@ -0,0 +1,75 @@
+from module.base.button import Button, ButtonWrapper
+
+# This file was auto-generated, do not modify it manually. To generate:
+# ``` python -m dev_tools.button_extract ```
+
+ICON_SEARCH = ButtonWrapper(
+    name='ICON_SEARCH',
+    share=Button(
+        file='./assets/share/combat/stamina/status/ICON_SEARCH.png',
+        area=(568, 8, 1265, 66),
+        search=(548, 0, 1280, 86),
+        color=(71, 70, 100),
+        button=(568, 8, 1265, 66),
+    ),
+)
+IMMERSIFIER_ICON = ButtonWrapper(
+    name='IMMERSIFIER_ICON',
+    share=Button(
+        file='./assets/share/combat/stamina/status/IMMERSIFIER_ICON.png',
+        area=(1047, 26, 1066, 49),
+        search=(1027, 6, 1086, 69),
+        color=(138, 127, 117),
+        button=(1047, 26, 1066, 49),
+    ),
+)
+IMMERSIFIER_OCR = ButtonWrapper(
+    name='IMMERSIFIER_OCR',
+    share=Button(
+        file='./assets/share/combat/stamina/status/IMMERSIFIER_OCR.png',
+        area=(1049, 26, 1151, 48),
+        search=(1029, 6, 1171, 68),
+        color=(64, 61, 61),
+        button=(1049, 26, 1151, 48),
+    ),
+)
+RESERVED_ICON = ButtonWrapper(
+    name='RESERVED_ICON',
+    share=Button(
+        file='./assets/share/combat/stamina/status/RESERVED_ICON.png',
+        area=(895, 26, 916, 48),
+        search=(875, 6, 936, 68),
+        color=(155, 212, 215),
+        button=(895, 26, 916, 48),
+    ),
+)
+RESERVED_OCR = ButtonWrapper(
+    name='RESERVED_OCR',
+    share=Button(
+        file='./assets/share/combat/stamina/status/RESERVED_OCR.png',
+        area=(895, 26, 999, 48),
+        search=(875, 6, 1019, 68),
+        color=(50, 69, 83),
+        button=(895, 26, 999, 48),
+    ),
+)
+STAMINA_ICON = ButtonWrapper(
+    name='STAMINA_ICON',
+    share=Button(
+        file='./assets/share/combat/stamina/status/STAMINA_ICON.png',
+        area=(873, 26, 894, 48),
+        search=(853, 6, 914, 68),
+        color=(188, 180, 226),
+        button=(873, 26, 894, 48),
+    ),
+)
+STAMINA_OCR = ButtonWrapper(
+    name='STAMINA_OCR',
+    share=Button(
+        file='./assets/share/combat/stamina/status/STAMINA_OCR.png',
+        area=(873, 26, 1013, 48),
+        search=(853, 6, 1033, 68),
+        color=(66, 64, 88),
+        button=(873, 26, 1013, 48),
+    ),
+)

tasks/combat/combat.py

@@ -54,7 +54,7 @@ class Combat(CombatInteract, CombatPrepare, CombatState, CombatTeam, CombatSuppo
 
         # Check limits
         if self.config.stored.TrailblazePower.value < self.combat_wave_cost:
-            return self._try_get_more_trablaize_power(self.config.stored.TrailblazePower.value, self.combat_wave_cost)
+            return self._try_get_more_trablaize_power(self.combat_wave_cost)
         if self.combat_waves <= 0:
             logger.info('Combat wave limited, cannot continue combat')
             return False
@@ -225,7 +225,7 @@ class Combat(CombatInteract, CombatPrepare, CombatState, CombatTeam, CombatSuppo
                 logger.info(f'Current has {current}, combat costs {self.combat_wave_cost}, can run again')
                 return True
             else:
-                return self._try_get_more_trablaize_power(current, self.combat_wave_cost * self.combat_waves)
+                return self._try_get_more_trablaize_power(self.combat_wave_cost * self.combat_waves)
         elif self.combat_wave_cost <= 0:
             logger.info(f'Free combat, combat costs {self.combat_wave_cost}, can not run again')
             return False
@@ -234,20 +234,15 @@ class Combat(CombatInteract, CombatPrepare, CombatState, CombatTeam, CombatSuppo
                 logger.info(f'Current has {current}, combat costs {self.combat_wave_cost}, can run again')
                 return True
             else:
-                return self._try_get_more_trablaize_power(current, self.combat_wave_cost * self.combat_waves)
-
-    def _try_get_more_trablaize_power(self, current, cost):
-        if self.config.TrailblazePower_ExtractReservedTrailblazePower:
-            logger.info('Extract reserved trailblaze power to get more trailblaze power')
-            if self.extract_reserved_trailblaze_power(current):
-                self.combat_get_trailblaze_power()
-                self.get_interval_timer(COMBAT_EXIT).wait()
-        if self.config.TrailblazePower_UseFuel:
-            logger.info('Use fuel to get more trailblaze power')
-            if self.use_fuel(current):
-                self.combat_get_trailblaze_power()
-            self.get_interval_timer(COMBAT_AGAIN).wait()
+                return self._try_get_more_trablaize_power(self.combat_wave_cost * self.combat_waves)
 
+    def _try_get_more_trablaize_power(self, cost):
+        self.extract_stamina(
+            update=False,
+            use_reserved=self.config.TrailblazePower_ExtractReservedTrailblazePower,
+            use_fuel=self.config.TrailblazePower_UseFuel
+        )
+        current = self.config.stored.TrailblazePower.value
         if current >= cost:
             return True
         else:

tasks/combat/fuel.py

@@ -4,14 +4,16 @@ from module.base.utils import area_offset, crop
 from module.logger import logger
 from module.ocr.ocr import Digit
 from tasks.base.assets.assets_base_popup import GET_REWARD, POPUP_CANCEL, POPUP_CONFIRM
-from tasks.base.ui import UI
 from tasks.combat.assets.assets_combat_finish import COMBAT_AGAIN, COMBAT_EXIT
-from tasks.combat.assets.assets_combat_fuel import *
 from tasks.combat.assets.assets_combat_prepare import COMBAT_PREPARE
+from tasks.combat.assets.assets_combat_stamina_fuel import *
+from tasks.combat.assets.assets_combat_stamina_reserved import *
+from tasks.combat.assets.assets_combat_stamina_status import *
+from tasks.combat.stamina_status import StaminaStatus
 from tasks.item.slider import Slider
 
 
-class Fuel(UI):
+class Fuel(StaminaStatus):
     fuel_trailblaze_power = 60
 
     def _use_fuel_finish(self):
@@ -26,7 +28,7 @@ class Fuel(UI):
                 return True
         if self.appear(COMBAT_PREPARE):
             if self.image_color_count(COMBAT_PREPARE.button, color=(230, 230, 230), threshold=240, count=400):
-                logger.info(f'Use fuel finished at COMBAT_AGAIN')
+                logger.info(f'Use fuel finished at COMBAT_PREPARE')
                 return True
         return False
 
@@ -87,7 +89,7 @@ class Fuel(UI):
         timer = self.get_interval_timer(COMBAT_EXIT, interval=5, renew=True)
         timer.set_current(4.4)
 
-    def extract_reserved_trailblaze_power(self, current, skip_first_screenshot=True):
+    def extract_reserved_trailblaze_power(self, skip_first_screenshot=True):
         """
         Extract reserved trailblaze power from previous combat.
 
@@ -95,11 +97,8 @@ class Fuel(UI):
             bool: If extracted
         """
         logger.info('Extract reserved trailblaze power')
-        reserved = Digit(OCR_RESERVED_TRAILBLAZE_POWER).ocr_single_line(self.device.image)
-        if reserved <= 0:
-            logger.info('No reserved trailblaze power')
-            return False
 
+        RESERVED_ICON.load_search(ICON_SEARCH.area)
         self.interval_clear([POPUP_CONFIRM, POPUP_CANCEL, GET_REWARD])
         while 1:
             if skip_first_screenshot:
@@ -111,12 +110,14 @@ class Fuel(UI):
                 break
             if self.appear_then_click(EXTRACT_RESERVED_TRAILBLAZE_POWER):
                 continue
-            if self.appear_then_click(RESERVED_TRAILBLAZE_POWER_ENTRANCE):
+            if self.appear_then_click(RESERVED_ICON):
                 continue
 
-        count = min(reserved, self.config.stored.TrailblazePower.FIXED_TOTAL - current)
-        logger.info(f'Having {reserved} reserved, going to use {count}')
-        self.set_reserved_trailblaze_power(count, total=reserved)
+        # No need, amount will be set by game client
+        # count = min(reserved, self.config.stored.TrailblazePower.FIXED_TOTAL - current)
+        # logger.info(f'Having {reserved} reserved, going to use {count}')
+        # self.set_reserved_trailblaze_power(count, total=reserved)
+
         self._fuel_confirm()
         return True
 
@@ -160,6 +161,7 @@ class Fuel(UI):
 
         logger.info("Use Fuel")
 
+        STAMINA_ICON.load_search(ICON_SEARCH.area)
         timeout = Timer(1, count=3)
         has_fuel = False
         while 1:
@@ -182,7 +184,7 @@ class Fuel(UI):
             if self.appear_then_click(FUEL):
                 has_fuel = True
                 continue
-            if not self.appear(POPUP_CONFIRM) and self.appear_then_click(FUEL_ENTRANCE):
+            if not self.appear(POPUP_CONFIRM) and self.appear_then_click(STAMINA_ICON):
                 continue
 
         offset = FUEL_SELECTED.button_offset
@@ -215,3 +217,39 @@ class Fuel(UI):
         self.set_fuel_count(use)
         self._fuel_confirm()
         return True
+
+    def extract_stamina(self, update=True, use_reserved=True, use_fuel=False):
+        """
+        Args:
+            update:
+            use_reserved:
+            use_fuel:
+
+        Returns:
+            bool: If used
+        """
+        if not use_reserved and not use_fuel:
+            return False
+
+        logger.hr('Extract stamina', level=2)
+        logger.info(f'Extract stamina, reserved={use_reserved}, fuel={use_fuel}')
+        if update:
+            self.update_stamina_status()
+        used = False
+
+        if use_reserved:
+            if self.config.stored.Reserved.value <= 0:
+                logger.info('No reserved stamina')
+            else:
+                self.extract_reserved_trailblaze_power()
+                used = True
+                self.update_stamina_status()
+                self.get_interval_timer(COMBAT_AGAIN).wait()
+
+        if use_fuel:
+            self.use_fuel(current=self.config.stored.TrailblazePower.value)
+            used = True
+            self.update_stamina_status()
+            self.get_interval_timer(COMBAT_AGAIN).wait()
+
+        return used

tasks/combat/prepare.py

@@ -1,35 +1,20 @@
-import re
-
 import module.config.server as server
 from module.base.timer import Timer
 from module.base.utils import color_similar, get_color
 from module.logger import logger
-from module.ocr.ocr import Digit, DigitCounter
-from tasks.base.ui import UI
+from module.ocr.ocr import Digit
 from tasks.combat.assets.assets_combat_prepare import (
-    OCR_TRAILBLAZE_POWER,
     OCR_WAVE_COST,
     OCR_WAVE_COUNT,
     WAVE_MINUS,
-    WAVE_PLUS, WAVE_SLIDER
+    WAVE_PLUS,
+    WAVE_SLIDER
 )
+from tasks.combat.stamina_status import StaminaStatus
 from tasks.item.slider import Slider
 
 
-class TrailblazePowerOcr(DigitCounter):
-    def after_process(self, result):
-        result = super().after_process(result)
-        # The trailblaze power icon is recognized as 买
-        # OCR_TRAILBLAZE_POWER includes the icon because the length varies by value
-        result = re.sub(r'[买米装：（）]', '', result)
-        # 61240 -> 6/240
-        result = re.sub(r'1240$', '/240', result)
-        # 0*0/24 -> 0/240
-        result = re.sub(r'24$', '240', result)
-        return result
-
-
-class CombatPrepare(UI):
+class CombatPrepare(StaminaStatus):
     # Current combat waves,
     combat_waves = 1
     # Limit combat runs, 0 means no limit.
@@ -78,20 +63,20 @@ class CombatPrepare(UI):
             else:
                 self.device.screenshot()
 
-            current, _, total = TrailblazePowerOcr(OCR_TRAILBLAZE_POWER).ocr_single_line(self.device.image)
-            # Empty result
-            if total == 0:
+            data = self.update_stamina_status(image=self.device.image)
+            if data.stamina is None:
                 continue
             # Confirm if it is > 240, sometimes just OCR errors
-            if current > 240 and timeout.reached():
+            # if current > 240 and timeout.reached():
+            #     break
+            if expect_reduce and timeout.reached():
                 break
-            if expect_reduce and current >= before:
+            if expect_reduce and data.stamina >= before:
                 continue
-            if current <= 240:
+            if data.stamina <= 240:
                 break
 
-        self.config.stored.TrailblazePower.value = current
-        return current
+        return data.stamina
 
     def combat_get_wave_cost(self, skip_first_screenshot=True):
         """

tasks/combat/stamina_status.py

@@ -0,0 +1,161 @@
+import re
+
+from pydantic import BaseModel
+
+from module.base.timer import Timer
+from module.base.utils import crop
+from module.logger import logger
+from module.ocr.ocr import Digit, DigitCounter
+from tasks.base.ui import UI
+from tasks.combat.assets.assets_combat_stamina_status import *
+
+
+class StaminaOcr(DigitCounter):
+    def after_process(self, result):
+        result = super().after_process(result)
+        # The trailblaze power icon is recognized as 买
+        # OCR_TRAILBLAZE_POWER includes the icon because the length varies by value
+        result = re.sub(r'[买米装来：（）]', '', result)
+        # 61240 -> 6/240
+        result = re.sub(r'1240$', '/240', result)
+        # 0*0/24 -> 0/240
+        result = re.sub(r'24$', '240', result)
+        return result
+
+
+class ReservedOcr(Digit):
+    pass
+
+
+class ImmersifierOcr(DigitCounter):
+    pass
+
+
+class DataStaminaStatus(BaseModel):
+    stamina: int | None
+    reserved: int | None
+    immersifier: int | None
+
+
+class StaminaStatus(UI):
+    @staticmethod
+    def get_stamina_status(image) -> DataStaminaStatus:
+        """
+        Update trailblaze power, stored trailblaze power, immersifier
+
+        Returns:
+            int: Stamina, or None if stamina not displayed or error on OCR
+            int: Reserved stamina
+            int: Immersifier
+        """
+        for button in [STAMINA_ICON, RESERVED_ICON, IMMERSIFIER_ICON]:
+            button.load_search(ICON_SEARCH.area)
+
+        stamina = None
+        if STAMINA_ICON.match_template(image):
+            STAMINA_OCR.load_offset(STAMINA_ICON)
+            im = crop(image, STAMINA_OCR.button, copy=False)
+            stamina, _, total = StaminaOcr(STAMINA_OCR).ocr_single_line(im, direct_ocr=True)
+            if total > 240 or total == 0:
+                logger.warning(f'Unexpected stamina total: {total}')
+                stamina = None
+
+        reserved = None
+        if RESERVED_ICON.match_template(image):
+            RESERVED_OCR.load_offset(RESERVED_ICON)
+            im = crop(image, RESERVED_OCR.button, copy=False)
+            reserved = ReservedOcr(RESERVED_OCR).ocr_single_line(im, direct_ocr=True)
+            if reserved > 2400:
+                logger.warning(f'Unexpected reserved value: {reserved}')
+                reserved = None
+
+        immersifier = None
+        if IMMERSIFIER_ICON.match_template(image):
+            IMMERSIFIER_OCR.load_offset(IMMERSIFIER_ICON)
+            im = crop(image, IMMERSIFIER_OCR.button, copy=False)
+            immersifier, _, total = StaminaOcr(IMMERSIFIER_OCR).ocr_single_line(im, direct_ocr=True)
+            if total != 8:
+                logger.warning(f'Unexpected immersifier total: {total}')
+                immersifier = None
+
+        return DataStaminaStatus(
+            stamina=stamina,
+            reserved=reserved,
+            immersifier=immersifier,
+        )
+
+    def update_stamina_status(
+            self,
+            image=None,
+            skip_first_screenshot=True,
+            expect_stamina=False,
+            expect_reserved=False,
+            expect_immersifier=False,
+    ) -> DataStaminaStatus:
+        """
+        Update stamina status with retry
+
+        Args:
+            image: Detect given image only, no new screenshot will be taken
+                and all expect_* are considered False
+            skip_first_screenshot:
+            expect_stamina:
+                True to expect stamina exists, retry detect if it wasn't
+            expect_reserved:
+            expect_immersifier:
+
+        Pages:
+            in: page_guild, Survival_Index, Simulated_Universe
+                or page_rogue, LEVEL_CONFIRM
+                or rogue, REWARD_CLOSE
+        """
+        timeout = Timer(1, count=2).start()
+        if image is None:
+            image = self.device.image
+            use_cached_image = False
+        else:
+            skip_first_screenshot = True
+            use_cached_image = True
+
+        while 1:
+            if skip_first_screenshot:
+                skip_first_screenshot = False
+            else:
+                self.device.screenshot()
+                image = self.device.image
+
+            # Timeout
+            if timeout.reached():
+                logger.warning('dungeon_update_stamina() timeout')
+                return DataStaminaStatus(
+                    stamina=None,
+                    reserved=None,
+                    immersifier=None,
+                )
+
+            # Ocr
+            status = self.get_stamina_status(image)
+            valid = True
+            if expect_stamina and status.stamina is None:
+                valid = False
+            if expect_reserved and status.reserved is None:
+                valid = False
+            if expect_immersifier and status.immersifier is None:
+                valid = False
+            if status.stamina is None and status.reserved is None and status.immersifier is None:
+                logger.warning('update_stamina_status: No icon detected')
+                valid = False
+
+            # Write config
+            with self.config.multi_set():
+                if status.stamina is not None:
+                    self.config.stored.TrailblazePower.value = status.stamina
+                if status.reserved is not None:
+                    self.config.stored.Reserved.value = status.reserved
+                if status.immersifier is not None:
+                    self.config.stored.Immersifier.value = status.reserved
+
+            if use_cached_image or valid:
+                return status
+            else:
+                continue

tasks/dungeon/assets/assets_dungeon_stamina.py

@@ -41,16 +41,6 @@ AMOUNT_PLUS = ButtonWrapper(
         ),
     ],
 )
-ENTER_IMMERSIFIER = ButtonWrapper(
-    name='ENTER_IMMERSIFIER',
-    share=Button(
-        file='./assets/share/dungeon/stamina/ENTER_IMMERSIFIER.png',
-        area=(1047, 26, 1066, 49),
-        search=(1027, 6, 1086, 69),
-        color=(138, 127, 117),
-        button=(1049, 26, 1151, 48),
-    ),
-)
 IMMERSIFIER_CHECK = ButtonWrapper(
     name='IMMERSIFIER_CHECK',
     share=Button(

tasks/dungeon/stamina.py

@@ -3,6 +3,7 @@ from module.base.timer import Timer
 from module.logger import logger
 from module.ocr.ocr import Digit
 from tasks.base.page import page_guide
+from tasks.combat.assets.assets_combat_stamina_status import ICON_SEARCH, IMMERSIFIER_ICON
 from tasks.dungeon.assets.assets_dungeon_stamina import *
 from tasks.dungeon.keywords import KEYWORDS_DUNGEON_TAB
 from tasks.dungeon.ui import DungeonUI
@@ -16,6 +17,7 @@ class DungeonStamina(DungeonUI):
             out: IMMERSIFIER_CHECK
         """
         logger.info('Enter immersifier')
+        IMMERSIFIER_ICON.load_search(ICON_SEARCH.area)
         while 1:
             if skip_first_screenshot:
                 skip_first_screenshot = False
@@ -24,7 +26,7 @@ class DungeonStamina(DungeonUI):
 
             if self.appear(IMMERSIFIER_CHECK):
                 break
-            if self.appear_then_click(ENTER_IMMERSIFIER, interval=2):
+            if self.appear_then_click(IMMERSIFIER_ICON, interval=2):
                 continue
 
     def _immersifier_exit(self, skip_first_screenshot=True):
@@ -140,7 +142,7 @@ class DungeonStamina(DungeonUI):
         logger.hr('Immersifier store', level=2)
         logger.info(f'Max store: {max_store}')
         self.dungeon_tab_goto(KEYWORDS_DUNGEON_TAB.Survival_Index)
-        self.dungeon_update_stamina()
+        self.update_stamina_status()
         before = self.config.stored.Immersifier.value
 
         if self.config.stored.Immersifier.is_full():
@@ -159,7 +161,7 @@ class DungeonStamina(DungeonUI):
         self._immersifier_enter()
         self._item_amount_set(amount, ocr_button=OCR_IMMERSIFIER_AMOUNT)
         self._item_confirm()
-        self.dungeon_update_stamina()
+        self.update_stamina_status()
         diff = self.config.stored.Immersifier.value - before
         logger.info(f'Stored {diff} immersifiers')
         return diff

tasks/dungeon/state.py

@@ -1,14 +1,13 @@
 from datetime import timedelta
 
 from module.base.base import ModuleBase
-from module.base.timer import Timer
 from module.base.utils import crop
 from module.config.stored.classes import now
 from module.config.utils import DEFAULT_TIME, get_server_next_monday_update, get_server_next_update
 from module.logger import logger
 from module.ocr.ocr import DigitCounter
-from tasks.base.ui import UI
-from tasks.dungeon.assets.assets_dungeon_state import OCR_SIMUNI_POINT, OCR_SIMUNI_POINT_OFFSET, OCR_STAMINA
+from tasks.combat.stamina_status import StaminaStatus
+from tasks.dungeon.assets.assets_dungeon_state import OCR_SIMUNI_POINT, OCR_SIMUNI_POINT_OFFSET
 from tasks.dungeon.keywords import DungeonList
 
 
@@ -19,7 +18,7 @@ class OcrSimUniPoint(DigitCounter):
         return result
 
 
-class DungeonState(UI):
+class DungeonState(StaminaStatus):
     def dungeon_get_simuni_point(self, image=None) -> int:
         """
         Page:
@@ -48,67 +47,6 @@ class DungeonState(UI):
             logger.warning(f'Invalid SimulatedUniverse points: {value}/{total}')
             return 0
 
-    def dungeon_update_stamina(self, image=None, skip_first_screenshot=True):
-        """
-        Returns:
-            bool: If success
-
-        Pages:
-            in: page_guild, Survival_Index, Simulated_Universe
-                or page_rogue, LEVEL_CONFIRM
-                or rogue, REWARD_CLOSE
-        """
-        ocr = DigitCounter(OCR_STAMINA)
-        timeout = Timer(1, count=2).start()
-        if image is None:
-            image = self.device.image
-            use_cached_image = False
-        else:
-            skip_first_screenshot = True
-            use_cached_image = True
-
-        while 1:
-            if skip_first_screenshot:
-                skip_first_screenshot = False
-            else:
-                self.device.screenshot()
-                image = self.device.image
-
-            stamina = (0, 0, 0)
-            immersifier = (0, 0, 0)
-
-            if timeout.reached():
-                logger.warning('dungeon_update_stamina() timeout')
-                return False
-
-            for row in ocr.detect_and_ocr(image):
-                if row.ocr_text.isdigit():
-                    continue
-                if row.ocr_text == '+':
-                    continue
-                if not ocr.is_format_matched(row.ocr_text):
-                    continue
-                data = ocr.format_result(row.ocr_text)
-                if data[2] == self.config.stored.TrailblazePower.FIXED_TOTAL:
-                    stamina = data
-                if data[2] == self.config.stored.Immersifier.FIXED_TOTAL:
-                    immersifier = data
-
-            if stamina[2] > 0 and immersifier[2] > 0:
-                break
-            if use_cached_image:
-                logger.info('dungeon_update_stamina() ended')
-                return
-
-        stamina = stamina[0]
-        immersifier = immersifier[0]
-        logger.attr('TrailblazePower', stamina)
-        logger.attr('Imersifier', immersifier)
-        with self.config.multi_set():
-            self.config.stored.TrailblazePower.value = stamina
-            self.config.stored.Immersifier.value = immersifier
-        return True
-
     def dungeon_update_simuni(self):
         """
         Update rogue weekly points, stamina, immersifier
@@ -122,8 +60,8 @@ class DungeonState(UI):
         def func(image):
             logger.info('Update thread start')
             with self.config.multi_set():
-                self.dungeon_get_simuni_point(image)
-                self.dungeon_update_stamina(image)
+                # self.dungeon_get_simuni_point(image)
+                self.update_stamina_status(image)
 
         ModuleBase.worker.submit(func, self.device.image)
 

tasks/ornament/combat.py

@@ -130,7 +130,7 @@ class OrnamentCombat(Dungeon, RouteLoader, DungeonState):
 
         after = before
         for _ in range(3):
-            self.dungeon_update_stamina()
+            self.update_stamina_status()
             after = self.get_equivalent_stamina()
             if expect_reduce:
                 if before > after:

tasks/rogue/entry/entry.py

@@ -261,7 +261,7 @@ class RogueEntry(RouteBase, RogueRewardHandler, RoguePathHandler, DungeonRogueUI
                 if not self.image_color_count(LEVEL_CONFIRM, color=(223, 223, 225), threshold=240, count=50):
                     self.interval_clear(LEVEL_CONFIRM)
                     continue
-                self.dungeon_update_stamina()
+                self.update_stamina_status()
                 self.check_stop_condition()
                 self.device.click(LEVEL_CONFIRM)
                 continue

tasks/rogue/event/reward.py

@@ -54,7 +54,7 @@ class RogueReward(RogueUI, CombatInteract, DungeonState):
                 confirm.reset()
                 continue
             if self.appear(REWARD_CLOSE, interval=2):
-                self.dungeon_update_stamina()
+                self.update_stamina_status()
                 if not init:
                     initial_stamina = self.config.stored.TrailblazePower.value
                     initial_immersifier = self.config.stored.Immersifier.value