Wiizard/wiizard/views/formatDialog.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/>.
#
########################################################################


from PyQt5.QtWidgets import (
  QDialog,
  QLabel,
  QGridLayout,
  QComboBox,
  QCheckBox,
  QPushButton,
  QMessageBox
)

from PyQt5.QtCore import QVariant

import pywwt

import wiizard.globals as Global
import wiizard.const as Const
from wiizard.translations import _


"""

 ============= FORMAT =============
  Partition:    [partition select ]
  Recover  :    [X]
                [CANCEL] [FORMAT]
 ==================================
"""


class FormatDialog(QDialog):
  def __init__(self, parent=None):
    QDialog.__init__(self, parent=parent)

    self.setWindowTitle(_("Format"))

    layout = QGridLayout(self)

    label = QLabel(_("Partition:"), self)
    self.partitionCombo = QComboBox(self)
    self.partitionCombo.setPlaceholderText(_("Select partition"))
    layout.addWidget(label              , 0, 0, 1, 1)
    layout.addWidget(self.partitionCombo, 0, 1, 1, 2)

    self.recoverCheck = QCheckBox(_("Try to recover Wii games"), self)
    self.recoverCheck.setToolTip(_("Try to recover Wii games on a WBFS "
                                 "formatted partition"))
    self.recoverCheck.setEnabled(False)
    layout.addWidget(self.recoverCheck  , 1, 1, 1, 2)

    cancelButton      = QPushButton(_("Cancel"), self)
    self.formatButton = QPushButton(_("Format"), self)
    layout.addWidget(cancelButton       , 2, 1, 1, 1)
    layout.addWidget(self.formatButton  , 2, 2, 1, 1)

    Global.DevicesModel.deviceAdded.connect(self.__onDeviceAdded)
    Global.DevicesModel.deviceRemoved.connect(self.__onDeviceRemoved)

    self.refreshPartitions()
    self.formatButton.setEnabled(False)

    self.partitionCombo.currentIndexChanged.connect(self.__onIndexChange)
    self.formatButton.clicked.connect(self.__onFormatClicked)
    cancelButton.clicked.connect(self.reject)

  def refreshPartitions(self):
    for path in Global.DevicesModel.devices:
      device = Global.DevicesModel.devices[path]
      self.__addDevice(device)

  def __onIndexChange(self, index):
    # Only enabled format button when a partition is selected
    if index == -1:
      self.formatButton.setEnabled(False)
    else:
      self.formatButton.setEnabled(True)

    # Only enable recover checkbox on wbfs partitions
    currentDevicePath = self.partitionCombo.currentData()
    if not currentDevicePath:
      self.recoverCheck.setEnabled(False)
      return
    device = Global.DevicesModel.devices[currentDevicePath]
    if device in Global.DevicesModel.wbfsPartitions:
      self.recoverCheck.setEnabled(True)
    else:
      self.recoverCheck.setEnabled(False)

  def __onFormatClicked(self):
    devicePath = self.partitionCombo.currentData()

    recover    = (self.recoverCheck.isEnabled() and self.recoverCheck.isChecked())
    recoverStr = ""
    if recover:
      recoverStr += (_("<br>However it will try to recover Wii games (no "
                     "guaranties!)."))

    answer = QMessageBox.question(self, _("Confirmation"),
                                  _("Are you sure you want to format "
                                 f"'<b>{devicePath}</b>' ?<br>"
                                  "<b>All data on that partition will be lost!</b>"
                                 f"{recoverStr}"))
    if answer != QMessageBox.Yes:
      return

    # Format the partition
    result = None
    try:
      result = pywwt.format_to_wbfs(devicePath, recover=recover, testMode=False)
    except pywwt.error as err:
      QMessageBox.critical(self, _("Failed"), _(f"Format failed! Error: {err}"))
    else:
      if result:
        QMessageBox.information(self, _("Success"), _("Format success!"))
        Global.DevicesModel.rescanWbfsPartitions()
      else:
        QMessageBox.critical(self, _("Failed"), _("Format failed!"))

    self.accept()

  def __addDevice(self, device):
    if device.type != "partition":
      return

    if device.isMounted():
      return

    fsType = device.getFsType()

    if fsType == "swap":
      return

    if not fsType:
      if device in Global.DevicesModel.wbfsPartitions:
        fsType = "wbfs"
      else:
        fsType = "unknown"

    size = device.getSize()
    # TODO make a format size function
    sizeLabel = "Bytes"
    if size > Const.GiB:
      size = round(size / Const.GiB, 2)
      sizeLabel = "GiB"
    elif size > Const.MiB:
      size = round(size / Const.MiB, 2)
      sizeLabel = "MiB"

    self.partitionCombo.addItem(f"{device.path} [{size} {sizeLabel}] "
                                f"[{device.vendor} - {device.model}] "
                                f"({fsType})", QVariant(device.path))

  def __onDeviceAdded(self, devicePath):
    device = Global.DevicesModel.devices[devicePath]
    self.__addDevice(device)

  def __onDeviceRemoved(self, devicePath):
    index = self.partitionCombo.findData(QVariant(devicePath))
    if index != -1:
      self.partitionCombo.removeItem(index)