import json import logging import time from arcaea_offline.database import Database from arcaea_offline.external.andreal.api_data import ( AndrealImageGeneratorApiDataConverter, ) from arcaea_offline.models import Chart from PySide6.QtCore import QCoreApplication, QDir, QFileInfo, Qt, Slot from PySide6.QtGui import QGuiApplication, QImage, QPainter, QPaintEvent, QPixmap from PySide6.QtWidgets import QButtonGroup, QFileDialog, QLabel, QMessageBox, QWidget from core.settings import SettingsKeys from ui.designer.tabs.tabTools.tabTools_Andreal_ui import Ui_TabTools_Andreal from ui.extends.shared.language import LanguageChangeEventFilter from ui.extends.tabs.tabTools.tabTools_Andreal import AndrealHelper from ui.implements.components.chartSelector import ChartSelector from ui.implements.components.songIdSelector import SongIdSelectorMode logger = logging.getLogger(__name__) class PreviewLabel(QLabel): def __init__(self, parent=None): super().__init__(parent) self.setWindowFlag(Qt.WindowType.Window, True) def show(self): super().show() # center the window width = self.width() height = self.height() screen = QGuiApplication.primaryScreen() screenWidth = screen.size().width() screenHeight = screen.size().height() self.setGeometry( max(0, screenWidth / 2 - width / 2), max(0, screenHeight / 2 - height / 2), min(width, screenWidth), min(height, screenHeight), ) def paintEvent(self, e: QPaintEvent) -> None: size = self.size() painter = QPainter(self) scaledPixmap = self.pixmap().scaled( size, Qt.AspectRatioMode.KeepAspectRatio, Qt.TransformationMode.SmoothTransformation, ) x = (size.width() - scaledPixmap.width()) / 2 y = (size.height() - scaledPixmap.height()) / 2 painter.drawPixmap(x, y, scaledPixmap) class ChartSelectorDialog(ChartSelector): def __init__(self, parent=None): super().__init__(parent) self.setWindowFlag(Qt.WindowType.Dialog, True) self.setSongIdSelectorMode(SongIdSelectorMode.Chart) class TabTools_Andreal(Ui_TabTools_Andreal, QWidget): def __init__(self, parent=None): super().__init__(parent) self.setupUi(self) self.languageChangeEventFilter = LanguageChangeEventFilter(self) self.installEventFilter(self.languageChangeEventFilter) self.db = Database() self.andrealHelper = AndrealHelper(self) self.andrealFolderSelector.setMode( self.andrealFolderSelector.getExistingDirectory ) self.andrealFolderSelector.filesSelected.connect(self.setHelperPaths) self.andrealExecutableSelector.filesSelected.connect(self.setHelperPaths) self.andrealFolderSelector.connectSettings(SettingsKeys.Andreal.Folder) self.andrealExecutableSelector.connectSettings(SettingsKeys.Andreal.Executable) self.generatePreviewButton.clicked.connect(self.requestPreview) self.generateImageButton.clicked.connect(self.requestGenerate) self.infoChart: Chart | None = None self.previewJsonPath = None self.generateJsonPath = None self.generateImageFormat = None self.andrealHelper.error.connect(self.previewError) self.andrealHelper.ready.connect(self.previewReady) self.andrealHelper.finished.connect(self.previewFinished) self.andrealHelper.error.connect(self.generateError) self.andrealHelper.ready.connect(self.generateReady) self.andrealHelper.finished.connect(self.generateFinished) self.imageTypeButtonGroup = QButtonGroup(self) self.imageTypeButtonGroup.addButton(self.imageType_infoRadioButton, 0) self.imageTypeButtonGroup.addButton(self.imageType_bestRadioButton, 1) self.imageTypeButtonGroup.addButton(self.imageType_best30RadioButton, 2) self.imageFormatButtonGroup = QButtonGroup(self) self.imageFormatButtonGroup.addButton(self.imageFormat_jpgRadioButton, 0) self.imageFormatButtonGroup.addButton(self.imageFormat_pngRadioButton, 1) self.imageTypeButtonGroup.idToggled.connect(self.fillImageVersionComboBox) self.fillImageVersionComboBox() self.chartSelectorDialog = ChartSelectorDialog(self) self.chartSelectorDialog.valueChanged.connect(self.chartValueUpdated) self.chartSelectButton.clicked.connect(self.chartSelectorDialog.show) def setHelperPaths(self): if selectedFiles := self.andrealFolderSelector.selectedFiles(): self.andrealHelper.andrealFolder = selectedFiles[0] if selectedFiles := self.andrealExecutableSelector.selectedFiles(): self.andrealHelper.andrealExecutable = selectedFiles[0] def chartValueUpdated(self): chart = self.chartSelectorDialog.value() self.infoChart = chart if chart: self.chartSelectLabel.setText( f"{chart.title}({chart.song_id}), {chart.rating_class}" ) @Slot() def on_imageTypeWhatIsThisButton_clicked(self): QMessageBox.information( self, None, QCoreApplication.translate("TabTools_Andreal", "imageWhatIsThisDialog.description"), ) # fmt: skip def imageFormat(self): buttonId = self.imageFormatButtonGroup.checkedId() return ["jpg", "png"][buttonId] if buttonId > -1 else None def imageType(self): buttonId = self.imageTypeButtonGroup.checkedId() return ["info", "best", "best30"][buttonId] if buttonId > -1 else None def fillImageVersionComboBox(self): imageType = self.imageType() if not imageType: return self.imageVersionComboBox.clear() if imageType in ["info", "best"]: self.imageVersionComboBox.addItem("3", 3) self.imageVersionComboBox.addItem("2", 2) self.imageVersionComboBox.addItem("1", 1) elif imageType == "best30": self.imageVersionComboBox.addItem("2", 2) self.imageVersionComboBox.addItem("1", 1) def imageVersion(self): return self.imageVersionComboBox.currentData() def requestComplete(self) -> bool: if not self.imageType(): return False imageType = self.imageType() if imageType == "best" and not self.infoChart: return False return self.imageVersion() is not None def getAndrealArguments(self, jsonFile: str, *, preview: bool = False): if not self.requestComplete(): return arguments = [ str(self.imageType()), "--json-file", jsonFile, "--img-version", str(self.imageVersion()), ] if self.andrealFolderSelector.selectedFiles(): arguments.append("--path") arguments.append(self.andrealFolderSelector.selectedFiles()[0]) if preview: arguments.extend(["--img-format", "jpg", "--img-quality", "20"]) else: arguments.extend(["--img-format", self.imageFormat()]) if self.imageFormat() == "jpg": arguments.extend(["--img-quality", str(self.jpgQualitySpinBox.value())]) return arguments def getAndrealJsonContent(self): if not self.requestComplete(): return None imageType = self.imageType() if imageType == "best" and not self.infoChart: return jsonContentDict = {} try: with self.db.sessionmaker() as session: converter = AndrealImageGeneratorApiDataConverter(session) if imageType == "info": jsonContentDict = converter.user_info() elif imageType == "best": jsonContentDict = converter.user_best( self.infoChart.song_id, self.infoChart.rating_class ) elif imageType == "best30": jsonContentDict = converter.user_best30() except Exception as e: logger.exception("getAndrealJsonContent error") QMessageBox.critical(self, None, str(e)) return ( json.dumps(jsonContentDict, ensure_ascii=False) if jsonContentDict else None ) def getAndrealJsonFileName(self): if not self.requestComplete(): return None imageType = self.imageType() timestamp = int(time.time() * 1000) fileNameParts = ["andreal", imageType] if imageType == "best": fileNameParts.extend([self.infoChart.song_id, self.infoChart.rating_class]) fileNameParts.append(timestamp) fileNameParts = [str(i) for i in fileNameParts] fileName = "-".join(fileNameParts) return f"{fileName}.json" def getTempAndrealJsonPath(self): if fileName := self.getAndrealJsonFileName(): return QDir.temp().filePath(fileName) else: return None @Slot() def on_exportJsonButton_clicked(self): content = self.getAndrealJsonContent() fileName = self.getAndrealJsonFileName() if not content or not fileName: return saveFileName, _ = QFileDialog.getSaveFileName(self, None, fileName) if not saveFileName: return with open(saveFileName, "w", encoding="utf-8") as jf: jf.write(content) def requestGenerate(self): jsonPath = self.getTempAndrealJsonPath() jsonContent = self.getAndrealJsonContent() if not jsonPath or not jsonContent: return self.generateImageButton.setEnabled(False) self.generateJsonPath = jsonPath self.generateImageFormat = self.imageFormat() with open(jsonPath, "w", encoding="utf-8") as jf: jf.write(jsonContent) self.andrealHelper.request(jsonPath, self.getAndrealArguments(jsonPath)) def generateFinished(self): self.generateImageButton.setEnabled(True) def generateError(self, jsonPath: str, errorMsg: str): if jsonPath != self.generateJsonPath: return QMessageBox.critical(self, "Generate Error", errorMsg) def generateReady(self, jsonPath: str, imageBytes: bytes): if jsonPath != self.generateJsonPath: return if not imageBytes: QMessageBox.critical(self, "Generate Error", "Empty bytes received.") return qImage = QImage.fromData(imageBytes) filePathParts = jsonPath.split(".") filePathParts[-1] = self.generateImageFormat filePath = ".".join(filePathParts) fileName = QFileInfo(filePath).fileName() saveFileName, _ = QFileDialog.getSaveFileName(self, None, fileName) if not saveFileName: return qImage.save(saveFileName, self.generateImageFormat) def requestPreview(self): jsonPath = self.getTempAndrealJsonPath() jsonContent = self.getAndrealJsonContent() if not jsonPath or not jsonContent: return self.generatePreviewButton.setEnabled(False) self.previewJsonPath = jsonPath with open(jsonPath, "w", encoding="utf-8") as jf: jf.write(jsonContent) self.andrealHelper.request( jsonPath, self.getAndrealArguments(jsonPath, preview=True) ) def previewFinished(self): self.generatePreviewButton.setEnabled(True) def previewError(self, jsonPath: str, errorMsg: str): if jsonPath != self.previewJsonPath: return QMessageBox.critical(self, "Preview Error", errorMsg) def previewReady(self, jsonPath: str, imageBytes: bytes): if jsonPath != self.previewJsonPath: return if not imageBytes: QMessageBox.critical(self, "Preview Error", "Empty bytes received.") return qImage = QImage.fromData(imageBytes) filePathParts = jsonPath.split(".") filePathParts.pop() filePath = ".".join(filePathParts) fileName = QFileInfo(filePath).fileName() previewLabel = PreviewLabel(self) previewLabel.setPixmap(QPixmap.fromImage(qImage)) previewLabel.setWindowTitle(f"preview {fileName}") previewLabel.show()