From 0a32329664689dd47a77624a6d7cde47ae9f58ee Mon Sep 17 00:00:00 2001
From: Zebartin <16185081+Zebartin@users.noreply.github.com>
Date: Tue, 12 Sep 2023 18:36:24 +0800
Subject: [PATCH] Fix: Assignments missing groups (#89)

---
 module/ui/draggable_list.py           |  5 +++++
 tasks/assignment/keywords/__init__.py |  3 +++
 tasks/assignment/ui.py                | 20 ++++++++++++++++----
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/module/ui/draggable_list.py b/module/ui/draggable_list.py
index e77473946..9b378afb3 100644
--- a/module/ui/draggable_list.py
+++ b/module/ui/draggable_list.py
@@ -159,6 +159,7 @@ class DraggableList:
             return False
 
         logger.info(f'Insight row: {row}, index={row_index}')
+        last_buttons: set[OcrResultButton] = None
         while 1:
             if skip_first_screenshot:
                 skip_first_screenshot = False
@@ -181,6 +182,10 @@ class DraggableList:
             main.wait_until_stable(self.search_button, timer=Timer(
                 0, count=0), timeout=Timer(1.5, count=5))
             skip_first_screenshot = True
+            if last_buttons == set(self.cur_buttons):
+                logger.warning(f'No more rows in {self}')
+                return False
+            last_buttons = set(self.cur_buttons)
 
         return True
 
diff --git a/tasks/assignment/keywords/__init__.py b/tasks/assignment/keywords/__init__.py
index d307fb7ec..a2d030258 100644
--- a/tasks/assignment/keywords/__init__.py
+++ b/tasks/assignment/keywords/__init__.py
@@ -20,6 +20,8 @@ KEYWORDS_ASSIGNMENT_GROUP.Synthesis_Materials.entries = (
     KEYWORDS_ASSIGNMENT_ENTRY.Spring_of_Life,
     KEYWORDS_ASSIGNMENT_ENTRY.The_Land_of_Gold,
     KEYWORDS_ASSIGNMENT_ENTRY.The_Blossom_in_the_Storm,
+    KEYWORDS_ASSIGNMENT_ENTRY.Legend_of_the_Puppet_Master,
+    KEYWORDS_ASSIGNMENT_ENTRY.The_Wages_of_Humanity,
 )
 for group in (
         KEYWORDS_ASSIGNMENT_GROUP.Character_Materials,
@@ -27,4 +29,5 @@ for group in (
         KEYWORDS_ASSIGNMENT_GROUP.Synthesis_Materials,
 ):
     for entry in group.entries:
+        assert entry.group is None
         entry.group = group
diff --git a/tasks/assignment/ui.py b/tasks/assignment/ui.py
index b4585cbc3..5e7a336c8 100644
--- a/tasks/assignment/ui.py
+++ b/tasks/assignment/ui.py
@@ -4,6 +4,7 @@ from typing import Iterator
 
 from module.base.base import ModuleBase
 from module.base.timer import Timer
+from module.exception import ScriptError
 from module.logger import logger
 from module.ocr.ocr import DigitCounter, Ocr
 from module.ui.draggable_list import DraggableList
@@ -37,12 +38,13 @@ class AssignmentSwitch(Switch):
 
 class AssignmentOcr(Ocr):
     OCR_REPLACE = {
-        'ch': [
+        'cn': [
             (KEYWORDS_ASSIGNMENT_ENTRY.Winter_Soldiers.name, '[黑]冬的战士们'),
             (KEYWORDS_ASSIGNMENT_ENTRY.Born_to_Obey.name, '[牛]而服从'),
             (KEYWORDS_ASSIGNMENT_ENTRY.Root_Out_the_Turpitude.name,
              '根除恶[擎薯尊掌鞋]?'),
             (KEYWORDS_ASSIGNMENT_ENTRY.Akashic_Records.name, '阿[未][夏复]记录'),
+            (KEYWORDS_ASSIGNMENT_ENTRY.Legend_of_the_Puppet_Master.name, '^师传说'),
             (KEYWORDS_ASSIGNMENT_ENTRY.The_Wages_of_Humanity.name, '[赠]养人类'),
         ]
     }
@@ -121,8 +123,17 @@ class AssignmentUI(UI):
             self.device.screenshot()
             self.goto_entry(KEYWORDS_ASSIGNMENT_ENTRY.Nameless_Land_Nameless_People)
         """
-        self.goto_group(entry.group)
-        ASSIGNMENT_ENTRY_LIST.select_row(entry, self)
+        if entry.group is None:
+            err_msg = f'{entry} is not in any group, please inform developers if possible'
+            logger.warning(err_msg)
+            for group in self._iter_groups():
+                self.goto_group(group)
+                if ASSIGNMENT_ENTRY_LIST.select_row(entry, self):
+                    return
+            raise ScriptError(err_msg)
+        else:
+            self.goto_group(entry.group)
+            ASSIGNMENT_ENTRY_LIST.select_row(entry, self)
 
     def _wait_until_entry_loaded(self):
         skip_first_screenshot = True
@@ -144,7 +155,8 @@ class AssignmentUI(UI):
     @property
     def _limit_status(self) -> tuple[int, int, int]:
         self.device.screenshot()
-        current, remain, total = DigitCounter(OCR_ASSIGNMENT_LIMIT).ocr_single_line(self.device.image)
+        current, remain, total = DigitCounter(
+            OCR_ASSIGNMENT_LIMIT).ocr_single_line(self.device.image)
         if total and current <= total:
             logger.attr('Assignment', f'{current}/{total}')
             self.config.stored.Assignment.set(current, total)