#!/usr/bin/python2.7
# -*- coding: utf-8 -*-
"""
An advanced multi boot live USB creator...

This software is licenced under General Public Licence (GPL).
This program is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY expressed or implied, including the implied warranties of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.


Copyright © 2010-2015  Sundar
Author(s): Sundar and Ian Bruce
Contributor(s): LiQiong Lee, kbytesys, Narfinger and many others.
"""

import sys
import os
from os.path import expanduser
import platform
import shutil
import subprocess

from PyQt4 import QtGui
from PyQt4 import QtCore
sys.path.append('scripts')  #  The script directory/module was not found if I dont include this line.
path = os.path.abspath(__file__)
sys.path.append(os.path.dirname(path))
from scripts.gui.ui_multibootusb import Ui_Dialog
import scripts.admin as admin
from scripts.syslinux import Syslinux
from scripts.iso import ISO
import scripts.gen_fun as gen_fun
import scripts.persistence as persistence
from scripts.install_distro import InstallDistro
from scripts.update_cfg import UpdateCfgFile
from scripts.uninstall_distro import UnInstall
import scripts.qemu as qemu
import scripts.config as config
from scripts.usb import USB


class AppGui(qemu.AppGui, QtGui.QDialog, Ui_Dialog):
    """
    Main multibootusb GUI manipulation class.
    """
    def __init__(self):
        QtGui.QDialog.__init__(self)
        self.ui = Ui_Dialog()
        self.ui.setupUi(self)
        self.usb = USB()
        self.install = InstallDistro()
        self.mbusb_dir = gen_fun.mbusb_dir()
        self.syslinux = Syslinux()
        self.update_cfg_files = UpdateCfgFile()

        #  Main Tab
        self.ui.detect_usb.clicked.connect(self.onRefereshClick)
        self.ui.close.clicked.connect(self.on_close_Click)
        self.ui.browse_iso.clicked.connect(self.browse_iso)
        self.ui.comboBox.activated[str].connect(self.onComboChange)

        self.ui.create.clicked.connect(lambda: self.onCreateClick(str(self.ui.comboBox.currentText()).strip(),
                                                                str(self.ui.lineEdit.text()).strip()))
        self.ui.uninstall.clicked.connect(self.OnUninstallClick)

        #  Syslinux Tab
        self.ui.install_syslinux.clicked.connect(self.onInstall_syslinuxClick)
        self.ui.edit_syslinux.clicked.connect(self.onedit_syslinux)

        # QEMU Tab
        self.ui.browse_iso_qemu.clicked.connect(self.on_Qemu_Browse_iso_Click)
        self.ui.boot_iso_qemu.clicked.connect(self.on_Qemu_Boot_iso_Click)
        self.ui.boot_usb_qemu.clicked.connect(lambda: self.on_Qemu_Boot_usb_Click(config.user_password, str(self.ui.comboBox.currentText())))
        self.ui.tabWidget.removeTab(3)

        #  Update progressbar and status
        self.progress_thread_install = GuiInstallProgress()
        self.progress_thread_install.finished.connect(self.install_syslinux)
        self.progress_thread_install.update.connect(self.ui.progressBar.setValue)
        self.progress_thread_install.status.connect(self.ui.status.setText)

        self.progress_thread_uninstall = GuiUninstallProgress()
        self.progress_thread_uninstall.finished.connect(self.uninstall_sys_file_update)
        self.progress_thread_uninstall.update.connect(self.ui.progressBar.setValue)
        self.progress_thread_uninstall.status.connect(self.ui.status.setText)

        if not admin.isUserAdmin():
            if platform.system() == "Linux":
                config.user_password = admin.get_password()
            elif platform.system() == "Windows":
                admin.runAsAdmin()
                sys.exit(0)

        self.add_device()

    def add_device(self):
        """
        Adds list of available USB devices to GUI combobox.
        :return:
        """
        detected_device = self.usb.list_usb()
        if bool(detected_device):
            for device in detected_device:
                self.ui.comboBox.addItem(str(device))
                if self.ui.comboBox.currentText():
                    self.onComboChange()
        #filename = raw_input()

        gen_fun.prepare_mbusb_dir()

    def onComboChange(self):
        """
        Detects and updates GUI with populated USB device details.
        :return:
        """
        config.usb_disk = str(self.ui.comboBox.currentText())
        if bool(config.usb_disk):
            self.usb_disk_details = self.usb.get_usb(config.usb_disk)
            config.usb_mount = self.usb_disk_details.mount
            self.ui.usb_size_ttl.setText("Filesystem :: " + self.usb.get_usb(config.usb_disk).filesystem)
            self.ui.usb_size_avl.setText("Free Size :: " + self.usb.get_usb(config.usb_disk).free_size)
            self.ui.usb_label.setText("Label :: " + self.usb.get_usb(config.usb_disk).label)
            self.ui.usb_dev.setText("Total Size :: " + self.usb.get_usb(config.usb_disk).total_size)
            self.ui.usb_mount.setText("Mount :: " + self.usb.get_usb(config.usb_disk).mount)
            self.update_list_box(self.usb.get_usb(config.usb_disk).mount)
        else:
            print "No USB disk found..."

        print "Device status is " + self.usb.get_usb(config.usb_disk).mount
        if self.usb.get_usb(config.usb_disk).mount == "Not mounted.":
            QtGui.QMessageBox.information(self, "USB not mounted...", "USB deivce " + config.usb_disk + " is not "
                                                                "mounted. \nPlease mount and rerun multibootusb.")

    def update_list_box(self, usb_mount):
        """
        Updates listbox with installed distros on selected USB disk.
        :param usb_mount: Selected USB disk from combobox.
        :return:
        """
        sys_cfg_file = os.path.join(str(usb_mount), "multibootusb", "syslinux.cfg")
        if os.path.exists(sys_cfg_file):
            self.ui.listWidget.clear()
            for line in open(sys_cfg_file):
                if "#start " in line:
                    self.ui.listWidget.addItem(line[7:])
            self.ui.status.clear()
        else:
            print "syslinux.cfg does not exist for updating list widget."

    def onRefereshClick(self):
        """
        Calls function to detect USB devices.
        :return:
        """
        self.ui.comboBox.clear()
        self.add_device()

    def browse_iso(self):
        if str(self.ui.lineEdit.text()):
            self.ui.lineEdit.clear()
        config.iso_link = str(QtGui.QFileDialog.getOpenFileName(self, 'Select an iso...', os.path.expanduser('~'), 'ISO Files (*.iso)'))
        if config.iso_link:
            if platform.system() == "Windows":
                if "/" in config.iso_link:
                    config.iso_link = config.iso_link.strip().replace("/", "\\")
            self.ui.lineEdit.insert(str(config.iso_link))
        else:
            print ("File not selected...")

    def onCreateClick(self, usb_disk, iso_link):
        """
        Main function to create bootable USB disk.
        :param usb_disk: ComboBox text as detected USB disk.
        :param iso_link: LineEdit text as selected ISO link.
        :return:
        """
        self.iso = ISO(config.iso_link)
        #config.usb_disk = str(self.ui.comboBox.currentText())

        if not usb_disk:
            QtGui.QMessageBox.information(self, "No Device...",
                                          "No USB device found.\n\nInsert USB and use Refresh USB button to detect USB.")
        elif not self.ui.lineEdit.text():
            QtGui.QMessageBox.information(self, "No ISO...", "No ISO found.\n\nPlease use step 2 to choose an ISO.")
        else:
            if not self.iso.integrity() is True:
                QtGui.QMessageBox.information(self, "Integrity...", "ISO integrity failed.\n\nPlease check the downloaded ISO.")
            else:
                if os.path.exists(iso_link):
                    self.ui.lineEdit.clear()
                    gen_fun.clean_iso_cfg_ext_dir(os.path.join(self.mbusb_dir, "iso_cfg_ext_dir"))  # Need to be cleaned.
                                                                                    # Otherwise it keeps all old distro files.
                    self.iso.iso_extract_file(os.path.join(self.mbusb_dir, "iso_cfg_ext_dir"), ".cfg")
                    self.iso.iso_extract_file(os.path.join(self.mbusb_dir, "iso_cfg_ext_dir"), ".CFG")
                    config.distro = self.iso.distro(os.path.join(self.mbusb_dir, "iso_cfg_ext_dir"))
                    if config.distro:
                        print "Distro detected is " + config.distro
                        if not gen_fun.distro_install_dir_exist(iso_link, self.usb.get_usb(config.usb_disk).mount) is True:
                            persistence_option = persistence.persistence_distro(config.distro)
                            if persistence_option is not None:
                                persistence_size = gen_fun.persistence_size(persistence_option)
                                config.persistence = persistence_size
                                #iso_size = gen_fun.install_size(iso_link, persistence_size)
                                config.install_size = gen_fun.install_size(iso_link, persistence_size)
                            else:
                                persistence_size = 0  # Later on required for update_cfg() function to update config files.
                                config.persistence = persistence_size
                                #iso_size = (os.path.getsize(iso_link) / 1024)
                                config.install_size = (os.path.getsize(iso_link) / 1024)
                            print "Persistence choosen is " + str(persistence_size) + " MB"

                            if config.install_size >= self.usb.disk_usage(self.usb.get_usb(usb_disk).mount):
                                QtGui.QMessageBox.information(self, "No Space.", "No space available on  " + usb_disk)
                            else:
                                reply = QtGui.QMessageBox.question(self, 'Review selection...',
                                                           'Selected USB disk:: %s\n' % usb_disk +
                                                           'USB mount point:: %s\n' % self.usb.get_usb(usb_disk).mount +
                                                           'Selected distro:: %s\n\n' % self.iso.iso_name() +
                                                           'Would you like to proceed for installation?',
                                                           QtGui.QMessageBox.Yes, QtGui.QMessageBox.No)

                                if reply == QtGui.QMessageBox.Yes:
                                    self.progress_thread_install.start()
                        else:
                            QtGui.QMessageBox.information(self, 'Already Exist...',
                                                      os.path.basename(iso_link) + ' is already installed.')
                    else:
                        QtGui.QMessageBox.information(self, 'No support...',
                                                      'Sorry.\n' + os.path.basename(iso_link) +
                                                      ' is not supported at the moment\n'
                                                      'Please email this issue to feedback.multibootusb@gmail.com')

    def install_syslinux(self):
        """
        Function to install syslinux on distro directory and on selected USB disks.
        :return:
        """
        self.ui.status.setText(str("Installing Syslinux..."))
        self.syslinux.install_distro()
        self.syslinux.install_default()
        self.update_cfg_files.distro_cfg_files()
        self.update_list_box(self.usb.get_usb(config.usb_disk).mount)
        if sys.platform.startswith("linux"):
            self.ui.status.setText("Sync is in progress...")
            os.system('sync')
        self.ui.status.clear()
        QtGui.QMessageBox.information(self, 'Finished...', self.iso.iso_name() + ' has been successfully installed.')

    def OnUninstallClick(self):
        """
        Triggers a function to uninstall a selected distro.
        :return:
        """
        if self.ui.listWidget.currentItem() is None:
            print "Please select a distro from the list."
            QtGui.QMessageBox.information(self, 'No selection.', 'Please select a distro from the list.')
        else:
            uninstall_distro_name = str(self.ui.listWidget.currentItem().text()).strip()
            self.uninstall = UnInstall()
            config.uninstall_distro = uninstall_distro_name
            print config.uninstall_distro

            reply = QtGui.QMessageBox.question(self, "Review selection...",
                                       "Are you sure to uninstall " + uninstall_distro_name,
                                      QtGui.QMessageBox.Yes, QtGui.QMessageBox.No)

            if reply == QtGui.QMessageBox.Yes:

                if not os.path.exists(self.usb.get_usb(config.usb_disk).mount + "/multibootusb/" + uninstall_distro_name):
                    print "Distro install directory not found. Just updating syslinux.cfg file."
                    self.uninstall.update_sys_cfg_file()
                else:
                    self.progress_thread_uninstall.start()

    def uninstall_sys_file_update(self):
        """
        Function to remove and update uninstall distro text.
        :return:
        """
        self.uninstall = UnInstall()
        self.uninstall.update_sys_cfg_file()
        self.update_list_box(self.usb.get_usb(config.usb_disk).mount)
        if sys.platform.startswith("linux"):
            self.ui.status.setText("Sync is in progress...")
            os.system('sync')
        self.ui.status.clear()
        QtGui.QMessageBox.information(self, 'Uninstall Complete...',  config.uninstall_distro + ' has been successfully removed.')

    def onInstall_syslinuxClick(self):
        """
        Function to install syslinux/extlinux on selected USB disk.
        :return:
        """
        if platform.system() == "Linux" or platform.system() == "Windows":

            if self.ui.install_sys_all.isChecked() or self.ui.install_sys_only.isChecked():
                print "Installing default syslinux on ", config.usb_disk
                ret = self.syslinux.install_default()
                if ret is True:
                    if self.ui.install_sys_all.isChecked():
                        print "Copying multibootusb directory to " + self.usb.get_usb(config.usb_disk).mount
                        for dirpath, dirnames, filenames in os.walk(gen_fun.resource_path(os.path.join("tools", "multibootusb"))):
                            for f in filenames:
                                print "Copying " + f
                                shutil.copy(gen_fun.resource_path(os.path.join(dirpath, f)),
                                    os.path.join(self.usb.get_usb(config.usb_disk).mount, "multibootusb"))
                    QtGui.QMessageBox.information(self, 'Install Success...',
                                      'Syslinux installed successfully on ' + config.usb_disk)
                elif ret is False:
                    QtGui.QMessageBox.information(self, 'Install error...',
                                      'Sorry. Syslinux failed to install on ' + config.usb_disk)
            else:
                QtGui.QMessageBox.information(self, 'No selection...',
                                          'Please select one of the option from above.')

    def onedit_syslinux(self):
        """
        Function to edit main syslinux.cfg file.
        :return:
        """
        # Function to edit syslinux.cfg file on editors like gedit, notepad etc.
        # Suggest me more editor which can be included in to this function.
        sys_cfg_file = os.path.join(self.usb.get_usb(config.usb_disk).mount, "multibootusb", "syslinux.cfg")
        print "Locating " + sys_cfg_file
        if not os.path.exists(sys_cfg_file):
            print "syslinux.cfg file not found..."
            QtGui.QMessageBox.information(self, 'File not found...', 'Sorry. Unable to locate syslinux.cfg file.\n'
                                                        'You can only edit syslinux.cfg file generated by multibootusb.')
        else:
            if platform.system() == "Linux":
                for e in config.editors_linux:
                    if subprocess.call('which ' + e, shell=True) == 0:
                        print "Editor found is " + e
                        editor = e
                        break
            elif platform.system() == "Windows":
                for e in config.editors_win:
                    if not gen_fun.which(e) is None:
                        print "Editor found is " + e
                        editor = e
                        break
            if not editor:
                QtGui.QMessageBox.information(self, 'Editor not found...', 'Sorry. Installed editor is not supported by multibootusb\n'
                                                                           'Edit syslinux.cfg file manually.\n')
            else:
                try:
                    subprocess.Popen(editor + " " + sys_cfg_file, shell=True).pid
                except OSError:
                    QtGui.QMessageBox.warning(self, 'Error...',
                                  'Failed to open syslinux.cfg file.\n'
                                  'Edit syslinux.cfg file manually.\n')

    def on_close_Click(self):
        """
        Closes main GUI.
        :return:
        """
        self.close()

    def closeEvent(self, event):
        """
        To capture the main close event.
        :param event: Close event.
        :return:
        """
        reply = QtGui.QMessageBox.question(self, 'Exit MultiBootUSB...',
                                           "Do you really want to quit multibootusb?", QtGui.QMessageBox.Yes,
                                           QtGui.QMessageBox.No)
        if reply == QtGui.QMessageBox.Yes:
            print "Closing multibootusb..."
            event.accept()
        else:
            print "Close event cancelled."
            event.ignore()


class GuiInstallProgress(QtCore.QThread):
    """
    Update GUI thread during install.
    """
    update = QtCore.pyqtSignal(int)
    status = QtCore.pyqtSignal(str)
    finished = QtCore.pyqtSignal()

    def __init__(self):
        QtCore.QThread.__init__(self)
        self.install = InstallDistro()
        self.usb = USB()
        self.thread = GenericThread(self.install.install_progress)
        #self.thread = self.install.install_progress

    def __del__(self):
        self.wait()

    def run(self):
        self.thread.start()
        self.iso = ISO(config.iso_link)

        install_dir = os.path.join(config.usb_mount, "multibootusb", self.iso.iso_basename())

        while self.thread.isRunning():
            if config.status_text.strip():
                status_text = config.status_text.replace(install_dir + "/", "Extracting ")
            else:
                status_text = config.status_text
            self.update.emit(config.percentage)
            self.status.emit(status_text)
            if not self.thread.isFinished() and config.percentage == 100:
                config.status_text = ""
                self.status.emit("Please wait...")

        self.update.emit(100)
        self.update.emit(0)

        self.status.emit("Installing boot loader...")

        if self.thread.isFinished():
            config.status_text = ""
            self.finished.emit()

        print "Distro extraction completed..."

        return


class GuiUninstallProgress(QtCore.QThread):
    """
    Update GUI thread during uninstall.
    """
    update = QtCore.pyqtSignal(int)
    status = QtCore.pyqtSignal(str)
    finished = QtCore.pyqtSignal()

    def __init__(self):
        QtCore.QThread.__init__(self)
        self.uninstall = UnInstall()
        self.usb = USB()
        self.thread = GenericThread(self.uninstall.uninstall_progress)

    def __del__(self):
        self.wait()

    def run(self):
        self.thread.start()
        while self.thread.isRunning():
            #print config.percentage
            self.update.emit(config.percentage)
            self.status.emit(config.status_text)
            if not self.thread.isFinished() and config.percentage == 100:
                config.status_text = "Please wait..."
            #time.sleep(.1)
        self.update.emit(100)
        self.update.emit(0)
        config.percentage = 0
        self.status.emit("Updating syslinux.cfg file...")

        if self.thread.isFinished():
            config.status_text = ""
            self.finished.emit()

        print "Distro uninstall is completed..."

        return


class GenericThread(QtCore.QThread):

    def __init__(self, function, *args, **kwargs):
        QtCore.QThread.__init__(self)
        self.function = function
        self.args = args
        self.kwargs = kwargs

    def __del__(self):
        self.wait()

    def run(self):
        self.function(*self.args, **self.kwargs)
        return

app = QtGui.QApplication(sys.argv)
window = AppGui()
ui = Ui_Dialog()
window.show()
window.setWindowTitle("MultiBootUSB - " + open(gen_fun.resource_path(os.path.join("tools", "version.txt")), 'r').read().strip())
window.setWindowIcon(QtGui.QIcon(gen_fun.resource_path(os.path.join("tools", "multibootusb.png"))))
sys.exit(app.exec_())