#!/usr/bin/python3

import os
import subprocess
import sys

from PyQt6.QtGui import QIcon
from PyQt6.QtWidgets import QApplication, QWidget, QGridLayout, QMainWindow, QPushButton, QComboBox, QGroupBox, \
    QVBoxLayout
from python3_mos_pyqt_dialogs.pyqt6_dialogs_functions import alert

version = '2.0'


class USBRepairWindow(QMainWindow):
    # Весь этот кусок взят из Rosa Image Writer и переписан на Python :)
    # // Using /sys/bus/usb/devices directory contents for enumerating the USB devices
    # //
    # // Details:
    # // Take the devices which have <device>/bInterfaceClass contents set to "08" (storage device).
    # //
    # // 1. To get the user-friendly name we need to read the <manufacturer> and <product> files
    # //    of the parent device (the parent device is the one with less-specified name, e.g. "2-1" for "2-1:1.0").
    # //
    # // 2. The block device name can be found by searching the contents of the following subdirectory:
    # //      <device>/host*/target*/<scsi-device-name>/block/
    # //    where * is a placeholder, and <scsi-device-name> starts with the same substring that "target*" ends with.
    # //    For example, this path may look like follows:
    # //      /sys/bus/usb/devices/1-1:1.0/host4/target4:0:0/4:0:0:0/block/
    # //    This path contains the list of block devices by their names, e.g. sdc, which gives us /dev/sdc.
    # //
    # // 3. And, finally, for the device size we multiply .../block/sdX/size (the number of sectors) with
    # //    .../block/sdX/queue/logical_block_size (the sector size).
    #
    # // Start with enumerating all the USB devices
    def get_usb_device_list(self):
        res = []
        all_devices = os.listdir('/sys/bus/usb/devices')
        for device in all_devices:
            device_folder = os.listdir(os.path.join('/sys/bus/usb/devices', device))
            if 'bInterfaceClass' not in device_folder:
                continue
            if not [i for i in device_folder if i.startswith('host')]:
                continue
            else:
                host = [i for i in device_folder if i.startswith('host')][0]
            with open(os.path.join('/sys/bus/usb/devices', device, 'bInterfaceClass')) as inp:
                if inp.read().strip() != '08':
                    continue
            host_folder = os.listdir(os.path.join('/sys/bus/usb/devices', device, host))
            if not [i for i in host_folder if i.startswith('target')]:
                continue
            else:
                target = [i for i in host_folder if i.startswith('target')][0]
                digits = target.replace('target', '')
            target_folder = os.listdir(os.path.join('/sys/bus/usb/devices', device, host, target))
            if not [i for i in target_folder if i.startswith(digits)]:
                continue
            else:
                new_digits = [i for i in target_folder if i.startswith(digits)][0]
            digits_folder = os.listdir(os.path.join('/sys/bus/usb/devices', device, host, target, new_digits))
            if 'block' not in digits_folder:
                continue
            else:
                device_name = \
                    os.listdir(os.path.join('/sys/bus/usb/devices', device, host, target, new_digits, 'block'))[0]
                number_of_sectors = int(open(
                    os.path.join('/sys/bus/usb/devices', device, host, target, new_digits, 'block', device_name,
                                 'size')).read())
                size_of_sector = int(open(
                    os.path.join('/sys/bus/usb/devices', device, host, target, new_digits, 'block', device_name,
                                 'queue', 'logical_block_size')).read())
                vendor = open(os.path.join('/sys/bus/usb/devices', device, host, target, new_digits, 'vendor')).read()
                gb = number_of_sectors * size_of_sector / 2 ** 30
            res.append([device_name.strip(), vendor.strip(), gb])
        res.sort(key=lambda x: x[0])
        return res

    def repair_usb(self):
        dev_address = os.path.join('/dev', self.device_list[self.usb_list_widget.currentIndex()][0])
        repair_method = f'ntfsfix' if self.repair_method_list.currentIndex() == 0 else '/sbin/fsck -A -y'
        # Если на устройстве несколько разделов, применять операцию нужно к ним, а не к самому блочному устройству
        partitions = subprocess.run(f'ls {dev_address}*', shell=True, capture_output=True).stdout.decode().split()
        shell_command = f'umount {" ".join(partitions)}'
        print(f'Running {shell_command}...')
        subprocess.run(shell_command, shell=True)
        success, errors = [], []
        # разделы не обнаружены, восстанавливать нужно само устройство
        if len(partitions) == 1:
            shell_command = f'{repair_method} {dev_address}'
            print(f'Running {shell_command}...')
            run_command = subprocess.run(shell_command, shell=True, capture_output=True)
            returncode = run_command.returncode
            if returncode == 0:
                success.append(dev_address)
            else:
                errors.append(dev_address)
            output = [run_command.stdout.decode()]
        else:
            # если список разделов выглядит как /dev/sda, /dev/sda1, /dev/sda2, нужно убрать само блочное устройство
            if dev_address in partitions:
                partitions.remove(dev_address)
            returncode = 0
            output = []

            for partition in partitions:
                shell_command = f'{repair_method} {partition}'
                print(f'Running {shell_command}...')
                run_command = subprocess.run(shell_command, shell=True, capture_output=True)
                returncode += run_command.returncode
                if run_command.returncode == 0:
                    success.append(partition)
                else:
                    errors.append(partition)
                output.append(run_command.stdout.decode())
        output = '\n'.join(output)
        print(output)

        if returncode == 0:
            alert('Успешно.')
        else:
            if 'is mounted' in output:
                alert('Восстановление не удалось. Размонтируйте устройство и повторите попытку.')
            else:
                text = ''
                if success:
                    success = "\n".join(success)
                    text += f'Восстановлены разделы:\n{success}\n'
                if errors:
                    errors = '\n'.join(errors)
                    text += f'\nНе удалось восстановить разделы:\n{errors}'
                alert(text)

    def __init__(self):
        QMainWindow.__init__(self)

        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        grid = QGridLayout()
        central_widget.setLayout(grid)
        device_groupbox = QGroupBox('Устройства')
        device_groupbox_layout = QVBoxLayout()
        device_groupbox.setLayout(device_groupbox_layout)
        self.usb_list_widget = QComboBox()
        self.device_list = self.get_usb_device_list()
        self.usb_list_widget.addItems([f'{i[0]} ({i[1]} - {i[2]:.2f} GB)' for i in self.device_list])
        device_groupbox_layout.addWidget(self.usb_list_widget)
        grid.addWidget(device_groupbox, 0, 0)

        methods_groupbox = QGroupBox('Методы восстановления')
        methods_groupbox_layout = QVBoxLayout()
        methods_groupbox.setLayout(methods_groupbox_layout)
        self.repair_method_list = QComboBox()
        self.repair_method_list.addItems(['ntfsfix', 'fsck'])
        methods_groupbox_layout.addWidget(self.repair_method_list)
        grid.addWidget(methods_groupbox, 0, 1)

        button = QPushButton('Восстановить')
        button.clicked.connect(self.repair_usb)
        if not self.device_list:
            button.setDisabled(True)
        button.setStyleSheet('color: green;')
        grid.addWidget(button, 1, 0, 1, 2)

        self.setGeometry(50, 50, 600, 50)
        self.setWindowTitle(f"Восстановление USB накопителей {version}")
        self.setWindowIcon(QIcon(QIcon.fromTheme('mos-usb-repair')))
        self.show()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    app.setDesktopFileName('mos-usb-repair')
    window = USBRepairWindow()
    sys.exit(app.exec())
