from forms.about import AboutWidget

from ui.mainwindow import Ui_MainWindow

from PyQt5 import QtWidgets, QtCore, QtGui

from devices import get_storage_devices

from mount.udisks import MountManager

import tempfile

import os

import shutil

import stat

import floppyemu

from .format_dlgs import FormatManyDialog, FormatOneDialog

class CustomListModel(QtCore.QStringListModel):

    def __init__(self, model=()):

        self.model = model

        super().__init__([m[1] for m in self.model])

    def setModel(self, model):

        self.model = model

        self.setStringList([m[1] for m in self.model])

    def getItem(self, index):

        return self.model[index]

class CustomTableModel(QtCore.QAbstractTableModel):

    headers = ['Label', 'Size']

    def __init__(self, placeholder):

        super().__init__()

        self.placeholder = placeholder

        self.model = []

    def setModel(self, model):

        self.beginResetModel()

        self.model = model

        self.endResetModel()

    def data(self, index, role=QtCore.Qt.DisplayRole):

        model = self.model[index.row()]

        if role == QtCore.Qt.DisplayRole:

            if QtCore.QVariant(model[index.column()]):

                if index.column() == 1:

                    return QtCore.QVariant(str(round(model[index.column()] / 1024, 2)) + 'KB')

                else:

                    return QtCore.QVariant(model[index.column()])

            else:

                return QtCore.QVariant(self.placeholder)

    def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole):

        if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:

            return QtCore.QVariant(self.headers[section])

        return super().headerData(section, orientation, role)

    def rowCount(self, parent=None):

        return len(self.model)

    def columnCount(self, parent=None):

        return 2

class MainWindow(Ui_MainWindow, QtWidgets.QMainWindow):

    def __init__(self, root=None):

        super().__init__()

        self.setupUi(self)

        self.root = root if root else tempfile.mktemp(prefix='kusbff_')

        os.makedirs(self.root, exist_ok=True)

        self.devices_model = CustomListModel()

        self.devices_combobox.setModel(self.devices_model)

        self.refresh_disk_drives()

        self.table_model = CustomTableModel('(Not set)')

        self.tableView.setModel(self.table_model)

        self.refresh_image_list()

        self.action_refresh_devices.triggered.connect(self.refresh_disk_drives)

        self.devices_combobox.currentIndexChanged.connect(self.on_device_change)

        self.action_open_floppy.triggered.connect(self.open_directory)

        self.action_write_floppy.triggered.connect(self.save_directory)

        self.action_read_image.triggered.connect(self.read_image)

        self.action_write_image.triggered.connect(self.write_image)

        self.action_format_disk.triggered.connect(self.format_one)

        self.action_format_usb.triggered.connect(self.format_many)

        self.action_about.triggered.connect(self.show_about_dialog)

    def show_about_dialog(self):

        AboutWidget().exec_()

    def format_one(self):

        if self.devices_combobox.currentIndex() != -1:

            indexes = self.tableView.selectedIndexes()

            device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0])

            FormatOneDialog.show_dialog(indexes[0].row(), device, os.path.join(self.root, 'temp'))

            self.refresh_image_list()

    def format_many(self):

        if self.devices_combobox.currentIndex() != -1:

            device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0])

            FormatManyDialog.show_dialog(device, os.path.join(self.root, 'temp'))

            self.refresh_image_list()

    def refresh_disk_drives(self):

        self.devices_model.setModel(get_storage_devices(True, 'disk'))

        self.refresh_image_list()

    def clean(self):

        MountManager.get().dispose_all()

        if os.path.exists(self.root):

            shutil.rmtree(self.root)

    def on_device_change(self):

        # Unmount and remove all old images

        MountManager.get().dispose_all()

        if os.path.exists(os.path.join(self.root, 'temp')):

            for f in os.listdir(os.path.join(self.root, 'temp')):

                if os.path.isfile(os.path.join(self.root, 'temp', f)):

                    os.remove(os.path.join(self.root, 'temp', f))

                else:

                    shutil.rmtree(os.path.join(self.root, 'temp', f))

        os.makedirs(os.path.join(self.root, 'temp'), exist_ok=True)

        self.refresh_image_list()

    def refresh_image_list(self):

        try:

            if self.devices_combobox.currentIndex() != -1:

                device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0])

                m = floppyemu.list_images(device)

                self.table_model.setModel(m)

        except PermissionError:

            device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0])

            QtWidgets.QMessageBox.critical(

                self, 'Permission error',

                'You have no access to the device %s. Please restart program as user with proper rights' % device

            )

    def open_directory(self):

        try:

            if self.devices_combobox.currentIndex() != -1:

                device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0])

                indexes = self.tableView.selectedIndexes()

                for index in indexes:

                    image = os.path.join(self.root, 'temp/image{}.img'.format(index.row()))

                    if MountManager.get().is_mounted(image):

                        MountManager.get().unmount(image)

                    if os.path.exists(image):

                        os.remove(image)

                    floppyemu.read_image(device, image, index.row())

                    os.makedirs(os.path.join(self.root, 'folders/{}'.format(index.row())), exist_ok=True)

                    os.chmod(os.path.join(self.root, 'folders/{}'.format(index.row())),

                             stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)

                    folder_path = MountManager.get().mount(image)

                    QtGui.QDesktopServices.openUrl(

                        QtCore.QUrl.fromLocalFile(os.path.join(self.root, folder_path)))

        except:

            import traceback

            traceback.print_exc()

    def save_directory(self):

        if self.devices_combobox.currentIndex() != -1:

            device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0])

            indexes = self.tableView.selectedIndexes()

            for index in indexes:

                image = os.path.join(self.root, 'temp/image{}.img'.format(index.row()))

                if MountManager.get().is_mounted(image):

                    MountManager.get().unmount(image)

                    floppyemu.write_image(device, image, index.row())

                    shutil.rmtree(os.path.join(self.root, 'folders/{}'.format(index.row())))

                else:

                    QtWidgets.QMessageBox.critical(

                        self, 'Error',

                        'Directory does not exist'

                    )

        self.refresh_image_list()

    def write_image(self):

        if self.devices_combobox.currentIndex() != -1:

            device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0])

            indexes = self.tableView.selectedIndexes()

            for index in indexes:

                fd = QtWidgets.QFileDialog(self)

                fd.setAcceptMode(QtWidgets.QFileDialog.AcceptOpen)

                fd.setWindowTitle('Select file for image {}'.format(index.row()))

                if fd.exec() == QtWidgets.QDialog.Accepted:

                    image = fd.selectedFiles()[0]

                    floppyemu.write_image(device, image, index.row())

                else:

                    return

        self.refresh_image_list()

    def read_image(self):

        if self.devices_combobox.currentIndex() != -1:

            device = '/dev/{}'.format(self.devices_model.getItem(self.devices_combobox.currentIndex())[0])

            indexes = self.tableView.selectedIndexes()

            for index in indexes:

                fd = QtWidgets.QFileDialog(self)

                fd.setAcceptMode(QtWidgets.QFileDialog.AcceptSave)

                fd.setWindowTitle('Select file for image {}'.format(index.row()))

                if fd.exec() == QtWidgets.QDialog.Accepted:

                    image = fd.selectedFiles()[0]

                    floppyemu.read_image(device, image, index.row())

                else:

                    return

    def closeEvent(self, event: QtGui.QCloseEvent):

        try:

            self.clean()

        except Exception as e:

            event.ignore()

            QtWidgets.QMessageBox.critical(

                self, 'Error',

                'Unable to clean up. Some resources are in use: {}'.format(e)

            )

<?xml version="1.0" encoding="UTF-8"?>

<ui version="4.0">

 <class>Dialog</class>

 <widget class="QDialog" name="Dialog">

  <property name="geometry">

   <rect>

    <x>0</x>

    <y>0</y>

    <width>400</width>

    <height>96</height>

   </rect>

  </property>

  <property name="windowTitle">

   <string>Format</string>

  </property>

  <layout class="QVBoxLayout" name="verticalLayout">

   <property name="spacing">

    <number>0</number>

   </property>

   <property name="leftMargin">

    <number>0</number>

   </property>

   <property name="topMargin">

    <number>0</number>

   </property>

   <property name="rightMargin">

    <number>0</number>

   </property>

   <property name="bottomMargin">

    <number>0</number>

   </property>

   <item>

    <layout class="QGridLayout" name="gridLayout">

     <item row="1" column="1">

     </item>

     <item row="0" column="1">

      <widget class="QLabel" name="label_2">

       <property name="maximumSize">

        <size>

         <width>16777215</width>

         <height>16</height>

        </size>

       </property>

       <property name="text">

        <string>to</string>

       </property>

      </widget>

     </item>

     <item row="1" column="0">

      <widget class="QSpinBox" name="sb_from"/>

     </item>

     <item row="0" column="0">

      <widget class="QLabel" name="label">

       <property name="maximumSize">

        <size>

         <width>16777215</width>

         <height>16</height>

        </size>

       </property>

       <property name="text">

        <string>from</string>

       </property>

      </widget>

     </item>

     <item row="2" column="0">

      <widget class="QRadioButton" name="rb_size_720">

       <property name="text">

        <string>&amp;720KB</string>

       </property>

      </widget>

     </item>

     <item row="2" column="1">

      <widget class="QRadioButton" name="rb_size_1440">

       <property name="text">

        <string>&amp;1.44MB</string>

       </property>

       <property name="checked">

        <bool>true</bool>

       </property>

      </widget>

     </item>

    </layout>

   </item>

   <item>

    <widget class="QDialogButtonBox" name="buttonBox">

     <property name="orientation">

      <enum>Qt::Horizontal</enum>

     </property>

     <property name="standardButtons">

      <set>QDialogButtonBox::Cancel|QDialogButtonBox::Ok</set>

     </property>

    </widget>

   </item>

  </layout>

 </widget>

 <resources/>

 <connections>

  <connection>

   <sender>buttonBox</sender>

   <signal>accepted()</signal>

   <receiver>Dialog</receiver>

   <slot>accept()</slot>

   <hints>

    <hint type="sourcelabel">

     <x>248</x>

     <y>254</y>

    </hint>

    <hint type="destinationlabel">

     <x>157</x>

     <y>274</y>

    </hint>

   </hints>

  </connection>

  <connection>

   <sender>buttonBox</sender>

   <signal>rejected()</signal>

   <receiver>Dialog</receiver>

   <slot>reject()</slot>

   <hints>

    <hint type="sourcelabel">

     <x>316</x>

     <y>260</y>

    </hint>

    <hint type="destinationlabel">

     <x>286</x>

     <y>274</y>

    </hint>

   </hints>

  </connection>

 </connections>

</ui>