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)

            )