From bbad8d03598afcd834947ff94434a231b89f86b6 Mon Sep 17 00:00:00 2001
From: LmeSzinc <37934724+LmeSzinc@users.noreply.github.com>
Date: Fri, 16 Aug 2024 22:31:22 +0800
Subject: [PATCH] Add: Handle story stages after A2

---
 campaign/event_20240815_cn/campaign_base.py | 75 +++++++++++++++++++++
 module/campaign/campaign_ocr.py             | 14 +++-
 module/campaign/campaign_ui.py              | 23 +++++--
 3 files changed, 102 insertions(+), 10 deletions(-)
 create mode 100644 campaign/event_20240815_cn/campaign_base.py

diff --git a/campaign/event_20240815_cn/campaign_base.py b/campaign/event_20240815_cn/campaign_base.py
new file mode 100644
index 000000000..959b5ddf4
--- /dev/null
+++ b/campaign/event_20240815_cn/campaign_base.py
@@ -0,0 +1,75 @@
+from module.base.timer import Timer
+from module.campaign.campaign_base import CampaignBase as CampaignBase_
+from module.combat.assets import GET_ITEMS_1
+from module.exception import CampaignNameError
+
+
+class CampaignBase(CampaignBase_):
+    entrance_timer = Timer(2)
+
+    def get_story_entrance(self):
+        """
+        Returns:
+            Button: Or None if nothing matched.
+        """
+        # 5 story stage after clearing A2
+        # You can't go anywhere unless you clicked it
+        return self.image_color_button(
+            area=(66, 200, 1200, 690), color=(0, 0, 0),
+            color_threshold=240, encourage=10, name='STORY_ENTRANCE')
+
+    def handle_story_entrance(self):
+        if not self.entrance_timer.reached():
+            return False
+
+        entrance = self.get_story_entrance()
+        if entrance is None:
+            return False
+
+        self.device.click(entrance)
+        self.entrance_timer.reset()
+        return True
+
+    def ensure_no_stage_entrance(self, skip_first_screenshot=True):
+        while 1:
+            if skip_first_screenshot:
+                skip_first_screenshot = False
+            else:
+                self.device.screenshot()
+
+            if self.is_in_stage_page():
+                # End
+                try:
+                    self._get_stage_name(self.device.image)
+                    return True
+                except (IndexError, CampaignNameError):
+                    pass
+                # Click
+                if self.handle_story_entrance():
+                    continue
+            if self.handle_story_skip():
+                self.interval_clear(GET_ITEMS_1)
+                self.entrance_timer.clear()
+                continue
+            if self.appear_then_click(GET_ITEMS_1, offset=(20, 20), interval=3):
+                self.entrance_timer.clear()
+                continue
+
+    def handle_in_stage(self):
+        # Click after stage ended
+        if self.is_in_stage_page():
+            if self.handle_story_entrance():
+                return True
+        return super().handle_in_stage()
+
+    def handle_get_chapter_additional(self):
+        # Exit when having story entrance
+        if self.get_story_entrance():
+            raise CampaignNameError
+        return super().handle_get_chapter_additional()
+
+    def handle_campaign_ui_additional(self):
+        if self.get_story_entrance():
+            self.ensure_no_stage_entrance()
+            return True
+        return super().handle_campaign_ui_additional()
diff --git a/module/campaign/campaign_ocr.py b/module/campaign/campaign_ocr.py
index b68219086..c414efe58 100644
--- a/module/campaign/campaign_ocr.py
+++ b/module/campaign/campaign_ocr.py
@@ -322,6 +322,15 @@ class CampaignOcr(ModuleBase):
         logger.attr('Chapter', self.campaign_chapter)
         logger.attr('Stage', ', '.join(self.stage_entrance.keys()))
 
+    def handle_get_chapter_additional(self):
+        """
+        Returns:
+            bool: If clicked
+        """
+        if self.appear(WITHDRAW, offset=(30, 30)):
+            logger.warning(f'get_chapter_index: WITHDRAW appears')
+            raise CampaignNameError
+
     def get_chapter_index(self, image):
         """
         A tricky method for ui_ensure_index
@@ -336,9 +345,8 @@ class CampaignOcr(ModuleBase):
         while 1:
             if timeout.reached():
                 raise CampaignNameError
-            if self.appear(WITHDRAW, offset=(30, 30)):
-                logger.warning(f'get_chapter_index: WITHDRAW appears')
-                raise CampaignNameError
+            if self.handle_get_chapter_additional():
+                continue
             try:
                 self._get_stage_name(image)
                 break
diff --git a/module/campaign/campaign_ui.py b/module/campaign/campaign_ui.py
index 75432d514..5912a5064 100644
--- a/module/campaign/campaign_ui.py
+++ b/module/campaign/campaign_ui.py
@@ -212,6 +212,21 @@ class CampaignUI(MapOperation, CampaignEvent, CampaignOcr):
         else:
             logger.warning(f'Unknown campaign chapter: {name}')
 
+    def handle_campaign_ui_additional(self):
+        """
+        Returns:
+            bool: If handled
+        """
+        if self.appear(WITHDRAW, offset=(30, 30)):
+            # logger.info("WITHDRAW button found, wait until map loaded to prevent bugs in game client")
+            self.ensure_no_info_bar(timeout=2)
+            try:
+                self.withdraw()
+            except CampaignEnd:
+                pass
+            return True
+        return False
+
     def ensure_campaign_ui(self, name, mode='normal'):
         for n in range(20):
             try:
@@ -220,13 +235,7 @@ class CampaignUI(MapOperation, CampaignEvent, CampaignOcr):
                 return True
             except CampaignNameError:
                 pass
-            if self.appear(WITHDRAW, offset=(30, 30)):
-                # logger.info("WITHDRAW button found, wait until map loaded to prevent bugs in game client")
-                self.ensure_no_info_bar(timeout=2)
-                try:
-                    self.withdraw()
-                except CampaignEnd:
-                    pass
+            if self.handle_campaign_ui_additional():
                 continue
 
             self.device.screenshot()