arcaea-offline-pyside-ui/ui/extends/shared/models/tables/score.py

import logging

from arcaea_offline.models import Chart, Difficulty, Score, ScoreCalculated, Song
from PySide6.QtCore import QCoreApplication, QModelIndex, QSortFilterProxyModel, Qt
from sqlalchemy import select

from .base import DbTableModel

logger = logging.getLogger(__name__)


class DbScoreTableModel(DbTableModel):
    IdRole = Qt.ItemDataRole.UserRole + 10
    ChartRole = Qt.ItemDataRole.UserRole + 11
    ScoreRole = Qt.ItemDataRole.UserRole + 12
    PttRole = Qt.ItemDataRole.UserRole + 13
    SongRole = Qt.ItemDataRole.UserRole + 14
    DifficultyRole = Qt.ItemDataRole.UserRole + 15

    def __init__(self, parent=None):
        super().__init__(parent)

        self.__items = []

    def retranslateHeaders(self):
        self._horizontalHeaders = [
            # fmt: off
            QCoreApplication.translate("DbScoreTableModel", "horizontalHeader.id"),
            QCoreApplication.translate("DbScoreTableModel", "horizontalHeader.chart"),
            QCoreApplication.translate("DbScoreTableModel", "horizontalHeader.score"),
            QCoreApplication.translate("DbScoreTableModel", "horizontalHeader.potential"),
            # fmt: on
        ]

    def syncDb(self):
        self.beginResetModel()
        self.beginRemoveRows(QModelIndex(), 0, self.rowCount())
        self.__items.clear()
        self.endRemoveRows()
        self.endResetModel()

        with self._db.sessionmaker() as session:
            stmt = (
                select(Score, Chart, Song, Difficulty, ScoreCalculated.potential)
                .join(
                    ScoreCalculated,
                    (Score.id == ScoreCalculated.id),
                    isouter=True,
                )
                .join(
                    Chart,
                    (Score.song_id == Chart.song_id)
                    & (Score.rating_class == Chart.rating_class),
                    isouter=True,
                )
                .join(
                    Song,
                    (Score.song_id == Song.id),
                    isouter=True,
                )
                .join(
                    Difficulty,
                    (Score.song_id == Difficulty.song_id)
                    & (Score.rating_class == Difficulty.rating_class),
                    isouter=True,
                )
            )
            results = session.execute(stmt).all()

            self.beginInsertRows(QModelIndex(), 0, len(results) - 1)
            for result in results:
                score, chart, song, difficulty, potential = result

                if chart:
                    chartInModel = chart
                elif song and difficulty:
                    chartInModel = Chart(
                        song_id=song.id,
                        rating_class=difficulty.rating_class,
                        title=difficulty.title or song.title,
                        set=song.set,
                    )
                else:
                    chartInModel = Chart(
                        song_id=score.song_id,
                        rating_class=score.rating_class,
                        title=score.song_id,
                        set="unknown",
                    )

                self.__items.append(
                    {
                        self.IdRole: score.id,
                        self.ScoreRole: score,
                        self.ChartRole: chartInModel,
                        self.SongRole: song,
                        self.DifficultyRole: difficulty,
                        self.PttRole: potential,
                    }
                )
            self.endInsertRows()

    def rowCount(self, *args):
        return len(self.__items)

    def data(self, index, role):
        if index.isValid() and self.checkIndex(index):
            if index.column() == 0 and role in [
                Qt.ItemDataRole.DisplayRole,
                self.IdRole,
            ]:
                return self.__items[index.row()][self.IdRole]
            elif index.column() == 1 and role in [
                self.ChartRole,
                self.SongRole,
                self.DifficultyRole,
            ]:
                return self.__items[index.row()][role]
            elif index.column() == 2 and role in [self.ChartRole, self.ScoreRole]:
                return self.__items[index.row()][role]
            elif index.column() == 3:
                potential = self.__items[index.row()][self.PttRole]
                if role == Qt.ItemDataRole.DisplayRole:
                    return f"{potential:.3f}" if potential is not None else "-"
                elif role == self.PttRole:
                    return potential
        return None

    def setData(self, index, value, role):
        if not (index.isValid() and self.checkIndex(index)):
            return False

        if index.column() == 2 and isinstance(value, Score) and role == self.ScoreRole:
            self._db.update_score(self.__items[index.row()][self.IdRole], value)
            self.syncDb()
            return True

        return False

    def flags(self, index) -> Qt.ItemFlag:
        flags = super().flags(index)
        flags |= Qt.ItemFlag.ItemIsSelectable
        if index.column() in [1, 2]:
            flags |= Qt.ItemFlag.ItemIsEditable
        return flags

    def _removeRow(self, row: int, syncDb: bool = True):
        if not 0 <= row < self.rowCount():
            return False

        try:
            self._db.delete_score(self.__items[row][self.ScoreRole])
            if syncDb:
                self.syncDb()
            return True
        except Exception:
            logger.exception(f"Table[Score]: Cannot remove row {row}")
            return False

    def removeRow(self, row: int, parent=...):
        return self._removeRow(row)

    def removeRows(self, row: int, count: int, parent=...):
        maxRow = min(self.rowCount() - 1, row + count - 1)
        if row > maxRow:
            return False

        result = all(
            self._removeRow(row, syncDb=False) for row in range(row, row + count)
        )
        self.syncDb()
        return result

    def removeRowList(self, rowList: list[int]):
        result = all(
            self._removeRow(row, syncDb=False) for row in sorted(rowList, reverse=True)
        )
        self.syncDb()
        return result


class DbScoreTableSortFilterProxyModel(QSortFilterProxyModel):
    Sort_C2_ScoreRole = Qt.ItemDataRole.UserRole + 75
    Sort_C2_TimeRole = Qt.ItemDataRole.UserRole + 76

    def lessThan(self, sourceLeft: QModelIndex, sourceRight: QModelIndex) -> bool:
        if sourceLeft.column() != sourceRight.column():
            return

        column = sourceLeft.column()
        if column == 0:
            return sourceLeft.data(DbScoreTableModel.IdRole) < sourceRight.data(
                DbScoreTableModel.IdRole
            )
        elif column == 2:
            scoreLeft = sourceLeft.data(DbScoreTableModel.ScoreRole)
            scoreRight = sourceRight.data(DbScoreTableModel.ScoreRole)
            if isinstance(scoreLeft, Score) and isinstance(scoreRight, Score):
                if self.sortRole() == self.Sort_C2_ScoreRole:
                    return scoreLeft.score < scoreRight.score
                elif self.sortRole() == self.Sort_C2_TimeRole:
                    if scoreLeft.date and scoreRight.date:
                        return scoreLeft.date < scoreRight.date
                    elif scoreLeft.date:
                        return False
                    else:
                        return True
        elif column == 3:
            pttLeft = sourceLeft.data(DbScoreTableModel.PttRole)
            pttRight = sourceRight.data(DbScoreTableModel.PttRole)
            if pttLeft and pttRight:
                return pttLeft < pttRight
            elif pttLeft:
                return False
            else:
                return True
        return super().lessThan(sourceLeft, sourceRight)