Add: Set combat waves

tasks/base/ui.py

@@ -2,6 +2,7 @@ from module.base.decorator import run_once
 from module.base.timer import Timer
 from module.exception import GameNotRunningError, GamePageUnknownError
 from module.logger import logger
+from module.ocr.ocr import Ocr
 from tasks.base.assets.assets_base_page import CLOSE
 from tasks.base.page import Page, page_main
 from tasks.base.popup import PopupHandler
@@ -152,6 +153,55 @@ class UI(PopupHandler):
             self.ui_goto(destination, skip_first_screenshot=True)
             return True
 
+    def ui_ensure_index(
+            self,
+            index,
+            letter,
+            next_button,
+            prev_button,
+            skip_first_screenshot=False,
+            fast=True,
+            interval=(0.2, 0.3),
+    ):
+        """
+        Args:
+            index (int):
+            letter (Ocr, callable): OCR button.
+            next_button (Button):
+            prev_button (Button):
+            skip_first_screenshot (bool):
+            fast (bool): Default true. False when index is not continuous.
+            interval (tuple, int, float): Seconds between two click.
+        """
+        logger.hr("UI ensure index")
+        retry = Timer(1, count=2)
+        while 1:
+            if skip_first_screenshot:
+                skip_first_screenshot = False
+            else:
+                self.device.screenshot()
+
+            if isinstance(letter, Ocr):
+                current = letter.ocr_single_line(self.device.image)
+            else:
+                current = letter(self.device.image)
+
+            logger.attr("Index", current)
+            diff = index - current
+            if diff == 0:
+                break
+            if current == 0:
+                logger.warning(f'ui_ensure_index got an empty current value: {current}')
+                continue
+
+            if retry.reached():
+                button = next_button if diff > 0 else prev_button
+                if fast:
+                    self.device.multi_click(button, n=abs(diff), interval=interval)
+                else:
+                    self.device.click(button)
+                retry.reset()
+
     def ui_goto_main(self):
         return self.ui_ensure(destination=page_main)
 

tasks/combat/assets/assets_combat_prepare.py

@@ -13,3 +13,33 @@ COMBAT_PREPARE = ButtonWrapper(
         button=(956, 640, 1224, 676),
     ),
 )
+OCR_WAVE_COUNT = ButtonWrapper(
+    name='OCR_WAVE_COUNT',
+    share=Button(
+        file='./assets/share/combat/prepare/OCR_WAVE_COUNT.png',
+        area=(911, 549, 1151, 581),
+        search=(891, 529, 1171, 601),
+        color=(27, 29, 32),
+        button=(911, 549, 1151, 581),
+    ),
+)
+WAVE_MINUS = ButtonWrapper(
+    name='WAVE_MINUS',
+    share=Button(
+        file='./assets/share/combat/prepare/WAVE_MINUS.png',
+        area=(825, 577, 858, 599),
+        search=(805, 557, 878, 619),
+        color=(239, 239, 239),
+        button=(825, 577, 858, 599),
+    ),
+)
+WAVE_PLUS = ButtonWrapper(
+    name='WAVE_PLUS',
+    share=Button(
+        file='./assets/share/combat/prepare/WAVE_PLUS.png',
+        area=(1203, 578, 1239, 598),
+        search=(1183, 558, 1259, 618),
+        color=(23, 24, 25),
+        button=(1203, 578, 1239, 598),
+    ),
+)

tasks/combat/prepare.py

@@ -0,0 +1,41 @@
+import re
+
+from module.logger import logger
+from module.ocr.ocr import Ocr
+from tasks.base.ui import UI
+from tasks.combat.assets.assets_combat_prepare import OCR_WAVE_COUNT, WAVE_MINUS, WAVE_PLUS
+
+
+class WaveCount(Ocr):
+    def after_process(self, result):
+        """
+        Returns:
+            int:
+        """
+        result = super().after_process(result)
+        logger.attr(name=self.name, text=str(result))
+
+        res = re.search(r'(\d)', result)
+        if res:
+            result = int(res.group(1))
+            if 1 <= result <= 6:
+                return result
+            else:
+                logger.warning(f'Unexpected combat wave: {result}')
+                return 0
+        else:
+            logger.warning('Cannot find wave count')
+            return 0
+
+
+class CombatPrepare(UI):
+    def combat_set_wave(self, count=6):
+        """
+        Args:
+            count: 1 to 6
+        """
+        self.ui_ensure_index(
+            count, letter=WaveCount(OCR_WAVE_COUNT),
+            next_button=WAVE_PLUS, prev_button=WAVE_MINUS,
+            skip_first_screenshot=True
+        )