Wiizard/wiizard/views/manageLocalSources.py

########################################################################
#  Wiizard - A Wii games manager
#  Copyright (C) 2023  CYBERDEViL
#
#  This file is part of Wiizard.
#
#  Wiizard is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  Wiizard is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
########################################################################


import os
import shutil
import math

from PyQt5.QtWidgets import (
  QWidget,
  QGridLayout,
  QLabel,
  QToolBar,
  QListWidget,
  QDialog,
  QComboBox,
  QSpinBox,
  QPushButton,
  QFileDialog,
  QMessageBox,
  QVBoxLayout,
  QTabWidget,
  QListWidgetItem,
  QApplication,
  QStyle
)
from PyQt5.QtCore import Qt, QVariant
from PyQt5.QtGui import QBrush

import pywwt

from wiizard import globals as Global
from wiizard import const as Const
from wiizard.models.source import LocalSourceModel, FilePartitionSourceModel


class AbstractManageLocalWidget(QWidget):
  def __init__(self, source, modelType, labelText, parent=None):
    QWidget.__init__(self, parent=parent)

    self.__newPaths  = []
    self.__delPaths  = []
    self._source     = source
    self.__modelType = modelType

    layout = QGridLayout(self)

    label = QLabel(labelText, self)
    label.setWordWrap(True)

    self.toolBar       = QToolBar(self)
    self.addAction     = self.toolBar.addAction("Add")
    self.delAction     = self.toolBar.addAction("Remove")
    self.restoreAction = self.toolBar.addAction("Restore")
    self.saveAction    = self.toolBar.addAction("Save")
    self.locationsList = QListWidget(self)

    layout.addWidget(label             , 0, 0)
    layout.addWidget(self.toolBar      , 1, 0)
    layout.addWidget(self.locationsList, 2, 0)

    # Add paths
    for path in self._source:
      sourceModel = self._source[path]
      item = QListWidgetItem(path, self.locationsList)
      if not sourceModel.isValid():
        item.setBackground(QBrush(Qt.darkRed))
        item.setIcon(QApplication.style().standardIcon(QStyle.SP_MessageBoxCritical))
        item.setToolTip("Missing")
      self.locationsList.addItem(item)

    self.__updateButtonStates()
    self.__selectionChanged()

    self.addAction.triggered.connect(self.__onAddPath)
    self.delAction.triggered.connect(self.__onDelPath)
    self.restoreAction.triggered.connect(self.restore)
    self.saveAction.triggered.connect(self.__onSave)
    self.locationsList.itemSelectionChanged.connect(self.__selectionChanged)

  def restore(self):
    self.__newPaths.clear()
    self.__delPaths.clear()
    self.locationsList.clear()
    for path in self._source:
      self.locationsList.addItem(path)
    self.__updateButtonStates()

  def hasUnsavedChanges(self):
    return (self.__newPaths or self.__delPaths)

  def __updateButtonStates(self):
    """ Update 'Save' and 'Restore' button states
    """
    if self.__newPaths or self.__delPaths:
      self.saveAction.setEnabled(True)
      self.restoreAction.setEnabled(True)
    else:
      self.saveAction.setEnabled(False)
      self.restoreAction.setEnabled(False)

  def __selectionChanged(self):
    """ Update 'Remove' button state
    """
    items = self.locationsList.selectedItems()
    if len(items) > 0:
      self.delAction.setEnabled(True)
    else:
      self.delAction.setEnabled(False)

  def addPath(self, path):
    if path:
      if path in self.__newPaths or path in self._source:
        if path in self.__delPaths and path in self._source:
          # In this scenario the path first got deleted by the user (so added
          # to __delPaths but still exists is _source), then the user adds the
          # same path again.
          self.__delPaths.remove(path)
          self.locationsList.addItem(path)
          self.__updateButtonStates()
        # Path already present ..
        return

      # Total new path
      self.__newPaths.append(path)
      self.locationsList.addItem(path)
      self.__updateButtonStates()

  def __onAddPath(self):
    path = self.getNewPath()
    self.addPath(path)

  def onDelPath(self, path):
    # For ManageLocalWbfsFilesWidget to be able to ask if the file should be
    # removed from disk.
    return True

  def save(self):
    self.__onSave()

  def __onDelPath(self):
    for item in self.locationsList.selectedItems():
      path = item.text()

      # Something gone wrong in the super class
      if not self.onDelPath(path):
        continue

      index = self.locationsList.row(item)
      self.locationsList.takeItem(index)

      if path in self.__newPaths:
        self.__newPaths.remove(path)
        continue

      self.__delPaths.append(path)
    self.__updateButtonStates()

  def __onSave(self):
    for path in self.__newPaths:
      model = self.__modelType(path)
      self._source.update({path: model})
    for path in self.__delPaths:
      del self._source[path]
    self.onSaved()
    self.parent().parent().parent().accept()  # this is ugly but it works

  def getNewPath(self):
    raise Exception("Re-implement this!")

  def onSaved(self):
    raise Exception("Re-implement this!")


class ManageLocalDirsWidget(AbstractManageLocalWidget):
  def __init__(self, parent=None):
    AbstractManageLocalWidget.__init__(self, Global.Sources.localDirectories,
                                       LocalSourceModel,
                                       "Add directories containing Wii games "
                                       "(.iso and .wbfs) below. It will "
                                       "recurse with a maximum depth of 4.",
                                       parent=parent)

  def getNewPath(self):
    return QFileDialog.getExistingDirectory(self, "Add Wii games directory")

  def onSaved(self):
    Global.Sources.localDirectoriesUpdated.emit()


class NewWbfsFileDialog(QDialog):
  def __init__(self, filepath, maxGiB, parent=None):
    QDialog.__init__(self, parent=parent)

    self.setWindowTitle("New WBFS file")

    layout = QGridLayout(self)

    label = QLabel("File:", self)
    value = QLabel(filepath, self)
    layout.addWidget(label        , 0, 0, 1, 1)
    layout.addWidget(value        , 0, 1, 1, 2)

    label = QLabel("Free space:", self)
    value = QLabel(f"{maxGiB} GiB", self)
    layout.addWidget(label        , 1, 0, 1, 1)
    layout.addWidget(value        , 1, 1, 1, 2)

    label = QLabel("Size:", self)
    self.sizeBox = QSpinBox(self)
    self.sizeBox.setSuffix(" GiB")
    self.sizeBox.setSingleStep(1)
    self.sizeBox.setMinimum(1)
    self.sizeBox.setMaximum(maxGiB)
    layout.addWidget(label        , 2, 0, 1, 1)
    layout.addWidget(self.sizeBox , 2, 1, 1, 2)

    label = QLabel("HD sector size:", self)
    self.hssCombo = QComboBox(self)
    self.hssCombo.addItem("auto"      , QVariant(0))
    self.hssCombo.addItem(" 512 bytes", QVariant(512))
    self.hssCombo.addItem("2048 bytes", QVariant(2048))
    self.hssCombo.addItem("4096 bytes", QVariant(4096))
    layout.addWidget(label        , 3, 0, 1, 1)
    layout.addWidget(self.hssCombo, 3, 1, 1, 2)

    label = QLabel("WBFS sector size:", self)
    self.wssCombo = QComboBox(self)
    self.wssCombo.addItem("auto"      , QVariant(0))
    self.wssCombo.addItem("1024 bytes", QVariant(1024))
    self.wssCombo.addItem("4096 bytes", QVariant(4096))
    self.wssCombo.addItem("8192 bytes", QVariant(8192))
    layout.addWidget(label        , 4, 0, 1, 1)
    layout.addWidget(self.wssCombo, 4, 1, 1, 2)

    cancelButton = QPushButton("Cancel", self)
    acceptButton = QPushButton("Accept", self)

    layout.addWidget(cancelButton, 5, 1, 1, 1)
    layout.addWidget(acceptButton, 5, 2, 1, 1)

    acceptButton.clicked.connect(self.accept)
    cancelButton.clicked.connect(self.reject)


class ManageLocalWbfsFilesWidget(AbstractManageLocalWidget):
  def __init__(self, parent=None):
    AbstractManageLocalWidget.__init__(self, Global.Sources.filePartitions,
                                       FilePartitionSourceModel,
                                       "Add WBFS partition files below.",
                                       parent=parent)
    self.newAction = self.toolBar.addAction("New")
    self.newAction.setToolTip("Create a new WBFS partition file.")

    self.newAction.triggered.connect(self.__onNewClicked)

  def getNewPath(self):
    output = QFileDialog.getOpenFileName(self, "Add WBFS partition file", "",
                                         "WBFS file (*.wbfs)")
    return output[0]

  def onSaved(self):
    Global.Sources.filePartitionsUpdated.emit()

  def onDelPath(self, path):
    source = self._source[path]

    # File doesn't exist any more, so no need to ask to delete the file.
    if not source.isValid():
      return True

    # Ask to delete the file
    result = QMessageBox.question(self, "Delete file?",
                                  f"Do you want to delete {path} from disk?")

    # Delete the file
    if result == QMessageBox.Yes:
      try:
        os.remove(path)
      except OSError as err:
        QMessageBox.critical(self, "Remove failed",
                             f"Removal of file {path} failed: {err}")
        return False

    # Don't delete the file
    return True

  def __onNewClicked(self):
    filepath = QFileDialog.getSaveFileName(self, "New WBFS partition file", "",
                                           "WBFS file (*.wbfs)")[0]

    if not filepath:
      return

    path    = os.path.dirname(filepath)
    freeGiB = int(math.floor(shutil.disk_usage(path).free / Const.GiB))

    if not freeGiB:
      QMessageBox.critical(self, "Not enough free space",
                           "Not enough free space at given location, the "
                           "minimum required free space is 1 GiB.")
      return

    if not filepath.endswith(".wbfs"):
      filepath += ".wbfs"

    sizeDialog = NewWbfsFileDialog(filepath, freeGiB, parent=self)
    if sizeDialog.exec():
      size = sizeDialog.sizeBox.value() 
      hss  = sizeDialog.hssCombo.currentData()
      wss  = sizeDialog.wssCombo.currentData()

      # TODO thread maybe the storage is slow?
      result = None
      try:
        result = pywwt.create_new_wbfs_file(filepath, size, hss=hss, wss=wss,
                                            testMode=False)
      except pywwt.error as err:
        QMessageBox.critical(self, "Failed!",
                             f"Creating new WBFS file failed! Error: {err}")
        return
      else:
        if not result:
          QMessageBox.critical(self, "Failed!", "Creating new WBFS file failed!")
          return

      self.addPath(filepath)


class ManageLocalWidget(QDialog):
  def __init__(self, parent=None):
    QDialog.__init__(self, parent=parent)
    self.setWindowTitle("Manage local sources")

    layout = QVBoxLayout(self)

    self.tabWidget = QTabWidget(self)

    self.localDirsWidget = ManageLocalDirsWidget(parent=self)
    self.localWbfsWidget = ManageLocalWbfsFilesWidget(parent=self)
    self.tabWidget.addTab(self.localDirsWidget, "Directories")
    self.tabWidget.addTab(self.localWbfsWidget, "WBFS Files")

    layout.addWidget(self.tabWidget)

  def closeEvent(self, closeEvent):
    if (self.localDirsWidget.hasUnsavedChanges()
        or self.localWbfsWidget.hasUnsavedChanges()):
      result = QMessageBox.question(self, "Save?",
                                    "There are unsaved changes, do you like "
                                    "to save them?",
                                    buttons=QMessageBox.Cancel|QMessageBox.No|QMessageBox.Yes)
      if result == QMessageBox.Cancel:
        closeEvent.ignore()
        return

      if result == QMessageBox.Yes:
        if self.localDirsWidget.hasUnsavedChanges():
          self.localDirsWidget.save()
        if self.localWbfsWidget.hasUnsavedChanges():
          self.localWbfsWidget.save()

      closeEvent.accept()