From afdef0371184756bb4ba9695d4f6efeabc1fd05b Mon Sep 17 00:00:00 2001
From: Zebartin <16185081+Zebartin@users.noreply.github.com>
Date: Tue, 26 Sep 2023 21:27:38 +0800
Subject: [PATCH] Fix: Check assignment status with timeout

---
 tasks/assignment/assignment.py | 21 +++++++++++----------
 tasks/assignment/ui.py         | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/tasks/assignment/assignment.py b/tasks/assignment/assignment.py
index a1f29bc3f..a3449dc4c 100644
--- a/tasks/assignment/assignment.py
+++ b/tasks/assignment/assignment.py
@@ -1,12 +1,10 @@
 from datetime import datetime
 
 from module.logger import logger
-from tasks.assignment.assets.assets_assignment_claim import CLAIM
-from tasks.assignment.assets.assets_assignment_dispatch import EMPTY_SLOT
-from tasks.assignment.assets.assets_assignment_ui import DISPATCHED
 from tasks.assignment.claim import AssignmentClaim
 from tasks.assignment.keywords import (KEYWORDS_ASSIGNMENT_GROUP,
                                        AssignmentEntry, AssignmentEventGroup)
+from tasks.assignment.ui import AssignmentStatus
 from tasks.base.page import page_assignment, page_menu
 from tasks.battle_pass.keywords import KEYWORD_BATTLE_PASS_QUEST
 from tasks.daily.keywords import KEYWORDS_DAILY_QUEST
@@ -105,13 +103,15 @@ class Assignment(AssignmentClaim, SynthesizeUI):
             logger.hr('Assignment inlist', level=2)
             logger.info(f'Check assignment inlist: {assignment}')
             self.goto_entry(assignment)
-            if self.appear(CLAIM):
+            status = self._check_assignment_status()
+            if status == AssignmentStatus.CLAIMABLE:
                 self.claim(assignment, duration, should_redispatch=True)
                 continue
-            if self.appear(DISPATCHED):
+            if status == AssignmentStatus.DISPATCHED:
                 self.dispatched[assignment] = datetime.now() + \
                     self._get_assignment_time()
                 continue
+            # General assignments must be dispatchable here
             if remain > 0:
                 self.dispatch(assignment, duration)
                 remain -= 1
@@ -138,11 +138,12 @@ class Assignment(AssignmentClaim, SynthesizeUI):
                 logger.hr('Assignment all', level=2)
                 logger.info(f'Check assignment all: {assignment}')
                 self.goto_entry(assignment)
-                if self.appear(CLAIM):
+                status = self._check_assignment_status()
+                if status == AssignmentStatus.CLAIMABLE:
                     self.claim(assignment, None, should_redispatch=False)
                     remain += 1
                     continue
-                if self.appear(DISPATCHED):
+                if status == AssignmentStatus.DISPATCHED:
                     self.dispatched[assignment] = datetime.now() + \
                         self._get_assignment_time()
                     if total == len(self.dispatched):
@@ -196,13 +197,13 @@ class Assignment(AssignmentClaim, SynthesizeUI):
                 logger.hr('Assignment event', level=2)
                 logger.info(f'Check assignment event: {assignment}')
                 self.goto_entry(assignment)
-                # No need to check dispatched here, should have been checked in _check_all
-                if self.appear(EMPTY_SLOT):
+                status = self._check_assignment_status()
+                # Should only be dispatchable or locked after _check_all
+                if status == AssignmentStatus.DISPATCHABLE:
                     self.dispatch(assignment, None)
                     remain -= 1
                     if remain <= 0:
                         return remain
                     continue
-                logger.info('Assignment is locked')
                 break
         return remain
diff --git a/tasks/assignment/ui.py b/tasks/assignment/ui.py
index 7fd7e2f0d..bb682a340 100644
--- a/tasks/assignment/ui.py
+++ b/tasks/assignment/ui.py
@@ -1,6 +1,7 @@
 import re
 from collections.abc import Iterator
 from datetime import timedelta
+from enum import Enum
 from functools import cached_property
 
 from module.base.timer import Timer
@@ -8,11 +9,20 @@ from module.exception import ScriptError
 from module.logger import logger
 from module.ocr.ocr import DigitCounter, Duration, Ocr
 from module.ui.draggable_list import DraggableList
+from tasks.assignment.assets.assets_assignment_claim import CLAIM
+from tasks.assignment.assets.assets_assignment_dispatch import EMPTY_SLOT
 from tasks.assignment.assets.assets_assignment_ui import *
 from tasks.assignment.keywords import *
 from tasks.base.ui import UI
 
 
+class AssignmentStatus(Enum):
+    CLAIMABLE = 0
+    DISPATCHED = 1
+    DISPATCHABLE = 2
+    LOCKED = 3
+
+
 class AssignmentOcr(Ocr):
     # EN has names in multiple rows
     merge_thres_y = 20
@@ -191,6 +201,28 @@ class AssignmentUI(UI):
             self.config.stored.Assignment.set(0, 0)
         return current, remain, total
 
+    def _check_assignment_status(self) -> AssignmentStatus:
+        skip_first_screenshot = True
+        timeout = Timer(2, count=3).start()
+        while 1:
+            if skip_first_screenshot:
+                skip_first_screenshot = False
+            else:
+                self.device.screenshot()
+
+            if timeout.reached():
+                logger.info(
+                    'Check assignment status timeout, assume LOCKED'
+                )
+                break
+            if self.appear(CLAIM):
+                return AssignmentStatus.CLAIMABLE
+            if self.appear(DISPATCHED):
+                return AssignmentStatus.DISPATCHED
+            if self.appear(EMPTY_SLOT):
+                return AssignmentStatus.DISPATCHABLE
+        return AssignmentStatus.LOCKED
+
     def _get_assignment_time(self) -> timedelta:
         return Duration(OCR_ASSIGNMENT_TIME).ocr_single_line(self.device.image)