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)
            )