From 716e872f0808efc6e511de9c657b9ded64ab92ec Mon Sep 17 00:00:00 2001
From: 283375 <bilibili283375@163.com>
Date: Sun, 17 Sep 2023 00:40:26 +0800
Subject: [PATCH] feat: `SongIdSelectorMode`

---
 ui/implements/components/chartSelector.py  | 11 ++++--
 ui/implements/components/songIdSelector.py | 41 ++++++++++++++++------
 2 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/ui/implements/components/chartSelector.py b/ui/implements/components/chartSelector.py
index b4f7c89..8a673fb 100644
--- a/ui/implements/components/chartSelector.py
+++ b/ui/implements/components/chartSelector.py
@@ -9,6 +9,7 @@ from PySide6.QtWidgets import QWidget
 from ui.designer.components.chartSelector_ui import Ui_ChartSelector
 from ui.extends.shared.database import databaseUpdateSignals
 from ui.extends.shared.language import LanguageChangeEventFilter
+from ui.implements.components.songIdSelector import SongIdSelectorMode
 
 logger = logging.getLogger(__name__)
 
@@ -36,6 +37,9 @@ class ChartSelector(Ui_ChartSelector, QWidget):
         )
         databaseUpdateSignals.songDataUpdated.connect(self.updateDatabase)
 
+    def setSongIdSelectorMode(self, mode: SongIdSelectorMode):
+        self.songIdSelector.setMode(mode)
+
     def value(self):
         songId = self.songIdSelector.songId()
         ratingClass = self.ratingClassSelector.value()
@@ -79,8 +83,11 @@ class ChartSelector(Ui_ChartSelector, QWidget):
         ratingClasses = []
         songId = self.songIdSelector.songId()
         if songId:
-            charts = self.db.get_charts_by_song_id(songId)
-            ratingClasses = [chart.rating_class for chart in charts]
+            if self.songIdSelector.mode == SongIdSelectorMode.Chart:
+                items = self.db.get_charts_by_song_id(songId)
+            else:
+                items = self.db.get_difficulties_by_song_id(songId)
+            ratingClasses = [item.rating_class for item in items]
         self.ratingClassSelector.setButtonsEnabled(ratingClasses)
 
     @Slot()
diff --git a/ui/implements/components/songIdSelector.py b/ui/implements/components/songIdSelector.py
index 1a229b0..0645cad 100644
--- a/ui/implements/components/songIdSelector.py
+++ b/ui/implements/components/songIdSelector.py
@@ -1,5 +1,6 @@
 import logging
 import re
+from enum import IntEnum
 from typing import Literal
 
 from arcaea_offline.database import Database
@@ -16,6 +17,11 @@ from ui.extends.shared.language import LanguageChangeEventFilter
 logger = logging.getLogger(__name__)
 
 
+class SongIdSelectorMode(IntEnum):
+    SongId = 0
+    Chart = 1
+
+
 class SongIdSelector(Ui_SongIdSelector, QWidget):
     valueChanged = Signal()
 
@@ -40,6 +46,8 @@ class SongIdSelector(Ui_SongIdSelector, QWidget):
             lambda: self.quickSwitchSelection("next", "package")
         )
 
+        self.mode = SongIdSelectorMode.SongId
+
         databaseUpdateSignals.songDataUpdated.connect(self.updateDatabase)
 
         self.fillPackComboBox()
@@ -62,6 +70,9 @@ class SongIdSelector(Ui_SongIdSelector, QWidget):
         self.packComboBox.currentIndexChanged.connect(self.valueChanged)
         self.songIdComboBox.currentIndexChanged.connect(self.valueChanged)
 
+    def setMode(self, mode: SongIdSelectorMode):
+        self.mode = mode
+
     def quickSwitchSelection(
         self,
         direction: Literal["previous", "next"],
@@ -135,20 +146,30 @@ class SongIdSelector(Ui_SongIdSelector, QWidget):
         self.songIdComboBox.clear()
         packId = self.packComboBox.currentData()
         if packId:
-            charts = self.db.get_charts_by_pack_id(packId)
-            inserted_song_ids = []
-            for chart in charts:
-                if chart.song_id not in inserted_song_ids:
-                    self.songIdComboBox.addItem(
-                        f"{chart.title} ({chart.song_id})", chart.song_id
-                    )
-                    inserted_song_ids.append(chart.song_id)
+            if self.mode == SongIdSelectorMode.SongId:
+                items = self.db.get_songs_by_pack_id(packId)
+            elif self.mode == SongIdSelectorMode.Chart:
+                items = self.db.get_charts_by_pack_id(packId)
+            else:
+                raise ValueError("Unknown SongIdSelectorMode.")
+            insertedSongIds = []
+            for item in items:
+                if self.mode == SongIdSelectorMode.SongId:
+                    itemId = item.id
+                elif self.mode == SongIdSelectorMode.Chart:
+                    itemId = item.song_id
+                else:
+                    continue
+
+                if itemId not in insertedSongIds:
+                    self.songIdComboBox.addItem(f"{item.title} ({itemId})", itemId)
+                    insertedSongIds.append(itemId)
                     row = self.songIdComboBox.count() - 1
                     self.songIdComboBox.setItemData(
-                        row, chart.title, DescriptionDelegate.MainTextRole
+                        row, item.title, DescriptionDelegate.MainTextRole
                     )
                     self.songIdComboBox.setItemData(
-                        row, chart.song_id, DescriptionDelegate.DescriptionTextRole
+                        row, itemId, DescriptionDelegate.DescriptionTextRole
                     )
         self.songIdComboBox.setCurrentIndex(-1)