Refactor: Abstract DungeonState class to be shared in rogue

tasks/dungeon/assets/assets_dungeon_state.py

@@ -0,0 +1,35 @@
+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 ```
+
+OCR_SIMUNI_POINT = ButtonWrapper(
+    name='OCR_SIMUNI_POINT',
+    share=Button(
+        file='./assets/share/dungeon/state/OCR_SIMUNI_POINT.png',
+        area=(580, 237, 860, 277),
+        search=(560, 217, 880, 297),
+        color=(163, 170, 252),
+        button=(580, 237, 860, 277),
+    ),
+)
+OCR_SIMUNI_POINT_OFFSET = ButtonWrapper(
+    name='OCR_SIMUNI_POINT_OFFSET',
+    share=Button(
+        file='./assets/share/dungeon/state/OCR_SIMUNI_POINT_OFFSET.png',
+        area=(685, 250, 717, 273),
+        search=(583, 187, 883, 387),
+        color=(199, 200, 250),
+        button=(685, 250, 717, 273),
+    ),
+)
+OCR_STAMINA = ButtonWrapper(
+    name='OCR_STAMINA',
+    share=Button(
+        file='./assets/share/dungeon/state/OCR_STAMINA.png',
+        area=(675, 11, 1181, 64),
+        search=(655, 0, 1201, 84),
+        color=(75, 89, 125),
+        button=(675, 11, 1181, 64),
+    ),
+)

tasks/dungeon/assets/assets_dungeon_ui.py

@@ -53,26 +53,6 @@ OCR_DUNGEON_NAV = ButtonWrapper(
         button=(117, 182, 423, 641),
     ),
 )
-OCR_SIMUNI_POINT = ButtonWrapper(
-    name='OCR_SIMUNI_POINT',
-    share=Button(
-        file='./assets/share/dungeon/ui/OCR_SIMUNI_POINT.png',
-        area=(580, 237, 820, 277),
-        search=(560, 217, 840, 297),
-        color=(166, 168, 252),
-        button=(580, 237, 820, 277),
-    ),
-)
-OCR_SIMUNI_POINT_OFFSET = ButtonWrapper(
-    name='OCR_SIMUNI_POINT_OFFSET',
-    share=Button(
-        file='./assets/share/dungeon/ui/OCR_SIMUNI_POINT_OFFSET.png',
-        area=(685, 250, 717, 273),
-        search=(583, 187, 883, 387),
-        color=(199, 200, 250),
-        button=(685, 250, 717, 273),
-    ),
-)
 OCR_WEEKLY_LIMIT = ButtonWrapper(
     name='OCR_WEEKLY_LIMIT',
     share=Button(

tasks/dungeon/state.py

@@ -0,0 +1,122 @@
+import threading
+
+from module.base.timer import Timer
+from module.base.utils import crop
+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
+
+
+class OcrSimUniPoint(DigitCounter):
+    def after_process(self, result):
+        result = super().after_process(result)
+        result = result.replace('O', '0').replace('o', '0')
+        return result
+
+
+class DungeonState(UI):
+    def dungeon_get_simuni_point(self, image=None) -> int:
+        """
+        Page:
+            in: page_guide, Survival_Index, Simulated_Universe
+        """
+        logger.info('Get simulated universe points')
+        if image is None:
+            image = self.device.image
+
+        _ = OCR_SIMUNI_POINT_OFFSET.match_template(image)
+        OCR_SIMUNI_POINT.load_offset(OCR_SIMUNI_POINT_OFFSET)
+        area = (
+            OCR_SIMUNI_POINT.area[0],
+            OCR_SIMUNI_POINT.button[1],
+            OCR_SIMUNI_POINT.area[2],
+            OCR_SIMUNI_POINT.button[3],
+        )
+
+        ocr = OcrSimUniPoint(OCR_SIMUNI_POINT)
+        value, _, total = ocr.ocr_single_line(crop(image, area), direct_ocr=True)
+        if total and value <= total:
+            logger.attr('SimulatedUniverse', f'{value}/{total}')
+            self.config.stored.SimulatedUniverse.set(value, total)
+            return value
+        else:
+            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
+        else:
+            skip_first_screenshot = True
+
+        while 1:
+            if skip_first_screenshot:
+                skip_first_screenshot = False
+            else:
+                self.device.screenshot()
+
+            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 in 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 image is not None:
+                logger.warning('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
+        Run in a new thread to be faster as data is not used immediately
+
+        Page:
+            in: page_guide, Survival_Index, Simulated_Universe
+        """
+        logger.info('dungeon_update_simuni')
+
+        def func(image):
+            logger.info('Update thread start')
+            with self.config.multi_set():
+                self.dungeon_get_simuni_point(image)
+                # self.dungeon_update_stamina(image)
+
+        thread = threading.Thread(target=func, args=(self.device.image,))
+        thread.start()

tasks/dungeon/ui.py

@@ -7,12 +7,11 @@ from module.base.button import ClickButton
 from module.base.timer import Timer
 from module.base.utils import get_color
 from module.logger import logger
-from module.ocr.ocr import DigitCounter, Ocr, OcrResultButton
+from module.ocr.ocr import Ocr, OcrResultButton
 from module.ocr.utils import split_and_pair_button_attr
 from module.ui.draggable_list import DraggableList
 from module.ui.switch import Switch
 from tasks.base.page import page_guide
-from tasks.base.ui import UI
 from tasks.combat.assets.assets_combat_prepare import COMBAT_PREPARE
 from tasks.dungeon.assets.assets_dungeon_ui import *
 from tasks.dungeon.keywords import (
@@ -24,6 +23,7 @@ from tasks.dungeon.keywords import (
     KEYWORDS_DUNGEON_TAB
 )
 from tasks.dungeon.keywords.classes import DungeonEntrance
+from tasks.dungeon.state import DungeonState
 
 
 class DungeonTabSwitch(Switch):
@@ -78,13 +78,6 @@ class OcrDungeonList(Ocr):
         return result
 
 
-class OcrSimUniPoint(DigitCounter):
-    def after_process(self, result):
-        result = super().after_process(result)
-        result = result.replace('O', '0').replace('o', '0')
-        return result
-
-
 class OcrDungeonListLimitEntrance(OcrDungeonList):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -123,7 +116,7 @@ DUNGEON_LIST = DraggableDungeonList(
     ocr_class=OcrDungeonList, search_button=OCR_DUNGEON_LIST)
 
 
-class DungeonUI(UI):
+class DungeonUI(DungeonState):
     def dungeon_tab_goto(self, state: DungeonTab):
         """
         Args:
@@ -233,31 +226,6 @@ class DungeonUI(UI):
                 logger.info('No Forgotten_Hall in list skip waiting')
                 return False
 
-    def dungeon_get_simuni_point(self) -> int:
-        """
-        Page:
-            in: page_guide, Survival_Index, Simulated_Universe
-        """
-        logger.info('Get simulated universe points')
-        _ = self.appear(OCR_SIMUNI_POINT_OFFSET)
-        OCR_SIMUNI_POINT.load_offset(OCR_SIMUNI_POINT_OFFSET)
-        area = (
-            OCR_SIMUNI_POINT.area[0],
-            OCR_SIMUNI_POINT.button[1],
-            OCR_SIMUNI_POINT.area[2],
-            OCR_SIMUNI_POINT.button[3],
-        )
-
-        ocr = OcrSimUniPoint(OCR_SIMUNI_POINT)
-        value, _, total = ocr.ocr_single_line(self.image_crop(area), direct_ocr=True)
-        if total and value <= total:
-            logger.attr('SimulatedUniverse', f'{value}/{total}')
-            self.config.stored.SimulatedUniverse.set(value, total)
-            return value
-        else:
-            logger.warning(f'Invalid SimulatedUniverse points: {value}/{total}')
-            return 0
-
     def _dungeon_nav_goto(self, dungeon: DungeonList, skip_first_screenshot=True):
         """
         Equivalent to `DUNGEON_NAV_LIST.select_row(dungeon.dungeon_nav, main=self)`
@@ -301,7 +269,7 @@ class DungeonUI(UI):
             logger.info('DUNGEON_NAV_LIST at top')
             # Update points if possible
             if DUNGEON_NAV_LIST.is_row_selected(button, main=self):
-                self.dungeon_get_simuni_point()
+                self.dungeon_update_simuni()
         else:
             # To start from any list states.
             logger.info('DUNGEON_NAV_LIST not at top')