#!/usr/bin/python3
import os
import subprocess
import sys

from PyQt5.QtGui import QIcon
from PyQt5.QtWidgets import QApplication, QWidget, QMainWindow, QPushButton, QVBoxLayout, QLabel, QHBoxLayout, \
    QComboBox, QScrollArea

from mos_live_resque_modules.design import question, alert


class MyWindow(QMainWindow):

    def get_partition_fs(self, partition):
        lines = subprocess.run(['blkid'], capture_output=True).stdout.decode().split('\n')
        for line in lines:
            if line.startswith(partition):
                return line.split('TYPE="')[1].split('"')[0]
        return ''

    def resque_grub(self):
        system_partition = self.partitions_combobox.currentText()
        if not question(message=f'Выбранный системный раздел: {system_partition}. Продолжить?'):
            return
        fs = self.get_partition_fs(system_partition)
        if not fs:
            alert(message='Не удалось определить тип ФС системного раздела! Восстановление невозможно.')
            return
        cmd = subprocess.run(['mkdir', '/chroot'], capture_output=True)
        print(cmd.stdout.decode())
        if fs == 'btrfs':
            cmd = subprocess.run(['mount', '-o', 'subvol=root', system_partition, '/chroot'], capture_output=True)
            print(cmd.stdout.decode())
        elif 'ext' in fs:
            cmd = subprocess.run(['mount', system_partition, '/chroot'], capture_output=True)
            print(cmd.stdout.decode())
        else:
            alert(f'Определённый тип файловой системы {fs} не является ни BTRFS, ни EXT. Восстановление невозможно.')
            return
        # cmd = subprocess.run(
        #     'for i in dev sys proc sys/firmware/efi/efivars; do mount --bind -v /$i /chroot/$i;done',
        #     shell=True, capture_output=True)
        for folder in ('dev', 'sys', 'proc', 'sys/firmware/efi/efivars'):
            cmd = subprocess.run(['mount', '--bind', '-v', f'/{folder}', f'/chroot/{folder}'], capture_output=True)
            print(cmd.stdout.decode())
        # print(cmd.stdout.decode())
        # cmd = subprocess.run(['chroot', '/chroot'], capture_output=True)
        os.chroot('/chroot')
        # print(cmd.stdout.decode())
        cmd = subprocess.run(['mount', '-a'], capture_output=True)
        print(cmd.stdout.decode())
        system_device = subprocess.run(
            f'lsblk -no PKNAME {system_partition}', shell=True, capture_output=True).stdout.decode().strip()
        print(f'{system_device=}')
        cmd = subprocess.run(['grub2-install', system_device], capture_output=True)
        print(cmd.stdout.decode())
        alert('Восстановление завершено.')

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

        if os.geteuid() != 0:
            alert('Данную программу нельзя запускать без повышенных привилегий.')
            sys.exit(0)

        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout()
        central_widget.setLayout(layout)

        scrollarea = QScrollArea()
        fdisk_output_label = QLabel()
        fdisk_output_label.setWordWrap(True)
        scrollarea.setWidget(fdisk_output_label)
        scrollarea.setWidgetResizable(True)
        layout.addWidget(scrollarea, stretch=2)
        fdisk_output_text = subprocess.run(['fdisk', '-l'], capture_output=True).stdout.decode()
        fdisk_output_label.setText(f'Вывод fdisk -l:\n{fdisk_output_text}')

        select_system_partition_layout = QHBoxLayout()
        layout.addLayout(select_system_partition_layout)
        select_system_partition_layout.addWidget(QLabel('Выберите системный раздел:'))
        self.partitions_combobox = QComboBox()
        partitions = []
        for line in fdisk_output_text.split('\n'):
            if line.strip().startswith('/dev'):
                last_word = line.split()[-1]
                if last_word.lower() != 'efi' and last_word.lower() != 'swap' and last_word.lower() != 'своп':
                    partitions.append(line.strip().split()[0])

        self.partitions_combobox.addItems(partitions)
        select_system_partition_layout.addWidget(self.partitions_combobox)

        layout.addStretch()

        button = QPushButton('Восстановить GRUB')
        button.clicked.connect(self.resque_grub)
        layout.addWidget(button)

        self.setGeometry(50, 50, 600, 500)
        self.setWindowTitle("Восстановление в live режиме")
        self.setWindowIcon(QIcon.fromTheme('mos-live-resque'))
        self.show()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MyWindow()
    sys.exit(app.exec_())
