#!/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): Joni Lee, kbytesys, Narfinger and many others.
"""

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

from PyQt4 import QtGui
from PyQt4 import QtCore

from scripts import usb
from scripts import admin
from scripts import iso
from scripts import syslinux
from scripts.gui.ui_multibootusb import Ui_Dialog
from scripts import gen_fun
from scripts import persistence
from scripts import update_cfg
from scripts import uninstall_distro
from scripts import update_status
from scripts import qemu
from scripts import config

user_password = ""
usb_disk = ""
usb_mount_path = ""
iso_link = ""
distro = ""
iso_size = ""
persistence_size = ""


class AppGui(uninstall_distro.AppGui, qemu.AppGui, QtGui.QDialog, Ui_Dialog):
    def __init__(self):
        QtGui.QDialog.__init__(self)
        global user_password
        self.ui = Ui_Dialog()
        self.ui.setupUi(self)

        #  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(),
                                       str(self.ui.usb_mount.text()[9:])))
        self.ui.uninstall.clicked.connect(lambda: self.uninstall_distro(str(self.ui.comboBox.currentText()).strip()))
        #  Syslinux Tab
        self.ui.install_syslinux.clicked.connect(self.onInstall_syslinuxClick)
        self.ui.edit_syslinux.clicked.connect(lambda: self.onedit_syslinux(str(self.ui.usb_mount.text()[9:])))

        # 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(user_password, str(self.ui.comboBox.currentText())))
        self.ui.tabWidget.removeTab(3)

        #  Update progressbar and status
        self.progress_thread_install = WorkThread()
        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)


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

        gen_fun.prepare_mbusb_dir()

    def browse_iso(self):
        global iso_link
        if str(self.ui.lineEdit.text()):
            self.ui.lineEdit.clear()
        iso_link = QtGui.QFileDialog.getOpenFileName(self, 'Select an iso...', expanduser("~"), 'ISO Files (*.iso)')
        if iso_link:
            self.ui.lineEdit.insert(iso_link)
        else:
            print ("File not selected.")

    def add_device(self):
        detected_device = usb.get_usb()
        if bool(detected_device):
            for device in detected_device:
                self.ui.comboBox.addItem(str(device))
                if self.ui.comboBox.currentText():
                    self.onComboChange()

    def onComboChange(self):
        global usb_disk
        usb_disk = str(self.ui.comboBox.currentText())
        usb_details = usb.usb_details(str(self.ui.comboBox.currentText()))
        if bool(usb_details):
            self.ui.usb_size_ttl.setText("Filesystem :: " + usb_details['filesystem'])
            self.ui.usb_size_avl.setText("Free Size :: " + usb_details['free_size'])
            self.ui.usb_label.setText("Label :: " + usb_details['label'])
            self.ui.usb_dev.setText("Total Size :: " + usb_details['total_size'])
            self.ui.usb_mount.setText("Mount :: " + usb_details['mount'])
            self.update_list_box(usb_details['mount'])
        print "Device ststus is " + usb_details['mount'].strip()
        if usb_details['mount'].strip() == "Not mounted.":
            QtGui.QMessageBox.information(self, "USB not mounted...", "USB deivce " + usb_disk + " is not mounted. \nPlease mount and rerun multibootusb.")

    def update_list_box(self, usb_mount):
        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:])
        else:
            print "syslinux.cfg does not exist for list update."

    def onRefereshClick(self):
        self.ui.comboBox.clear()
        self.add_device()

    def onCreateClick(self, usb_disk, iso_link, usb_mount):
        global usb_mount_path
        global distro
        global persistence_size
        global iso_size
        usb_mount_path = usb_mount

        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 os.path.exists(iso_link):
                self.ui.lineEdit.clear()
                home = gen_fun.mbusb_dir()
                gen_fun.clean_iso_cfg_ext_dir(os.path.join(home, "iso_cfg_ext_dir"))  # Need to be cleaned. Otherwise it keeps old distro files.
                iso.iso_extract_file(iso_link, os.path.join(home, "iso_cfg_ext_dir"), ".cfg")
                iso.iso_extract_file(iso_link, os.path.join(home, "iso_cfg_ext_dir"), ".CFG")
                distro = iso.detect_iso(iso_link, os.path.join(home, "iso_cfg_ext_dir"))
                if distro:
                    print "Distro detected is " + distro
                    if not gen_fun.distro_install_dir_exist(iso_link, usb_mount_path):
                        persistence_option = persistence.persistence_distro(distro)
                        if persistence_option is not None:
                            persistence_size = gen_fun.persistence_size(persistence_option)
                            iso_size = gen_fun.install_size(iso_link, persistence_size)
                        else:
                            persistence_size = 0  # Later on required for update_cfg() function to update config files.
                            iso_size = (os.path.getsize(iso_link) / 1024)
                        print "Persistence choosen is " + str(persistence_size) + " MB"

                        if iso_size >= usb.disk_usage(usb_mount_path):
                            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' % usb_mount_path +
                                                       'Selected distro:: %s\n\n' % iso.iso_name(iso_link) +
                                                       'Would you like to proceed for installation?',
                                                       QtGui.QMessageBox.Yes, QtGui.QMessageBox.No)

                            if reply == QtGui.QMessageBox.Yes:
                                self.update_status = update_status.UpdateStatusThread()
                                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):

        syslinux.install_distro_syslinux(distro, iso_link, usb_disk, user_password)
        syslinux.install_default_syslinux(user_password, usb_disk)
        update_cfg.update_distro_cfg_files(distro, str(self.ui.comboBox.currentText()).strip(),
                                           iso.iso_name(iso_link),
                                           gen_fun.install_dir(iso_link, usb_mount_path), persistence_size, iso_link)

        self.ui.status.setText("Updating config files...")
        self.update_list_box(usb_mount_path)
        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...', iso.iso_name(iso_link) + ' has been successfully installed.')

    def onInstall_syslinuxClick(self):

        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 ", usb_disk
                ret = syslinux.install_default_syslinux(user_password, usb_disk)
                if ret is True:
                    QtGui.QMessageBox.information(self, 'Install Success...',
                                      'Syslinux installed successfully on ' + usb_disk)
                elif ret is False:
                    QtGui.QMessageBox.information(self, 'Install error...',
                                      'Sorry. Syslinux failed to install on ' + usb_disk)
            elif self.ui.install_sys_all.isChecked():
                if not os.path.exists(os.path.join(usb_mount_path, "multibootusb")):
                    print "Copying multibootusb directory to " + usb_mount_path
                    shutil.copytree(gen_fun.resource_path(os.path.join("tools", "multibootusb")),
                                    os.path.join(usb_mount_path, "multibootusb"))
                    print "Installing default syslinux on ", usb_disk
                    ret = syslinux.install_default_syslinux(user_password, usb_disk)
                    if ret is True:
                        QtGui.QMessageBox.information(self, 'Install Success...',
                                          'Syslinux installed successfully on ' + usb_disk)
                    elif ret is False:
                        QtGui.QMessageBox.information(self, 'Install error...',
                                          'Sorry. Syslinux failed to install on ' + usb_disk)
            else:
                QtGui.QMessageBox.information(self, 'No selection...',
                                          'Please select one of the option from above.')

    def onedit_syslinux(self,usb_mount_path):
        # 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(usb_mount_path, "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):
        self.close()


    def closeEvent(self, event):
        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 WorkThread(QtCore.QThread):
    global usb_mount_path
    global persistence_size


    update = QtCore.pyqtSignal(int)
    status = QtCore.pyqtSignal(str)
    finished = QtCore.pyqtSignal()

    def __init__(self):
        QtCore.QThread.__init__(self)

    def __del__(self):
        self.wait()

    def run(self):
        install_dir = gen_fun.install_dir(iso_link, usb_mount_path)
        self.thread = GenericThread(iso.install_distro, distro, iso_link, usb_mount_path, install_dir, persistence_size)
        self.thread.start()
        self.update_thread = update_status.UpdateStatusThread()
        install_size = usb.disk_usage(usb_mount_path).used / 1024
        while self.thread.isRunning():
            current_size = usb.disk_usage(usb_mount_path).used / 1024
            diff_size = abs(int(current_size - install_size))
            percentage = float(1.0 * diff_size) / iso_size * 100
            text = config.iso_extract_file_name.replace(install_dir, "Extracting ")
            #print percentage
            self.update.emit(percentage)
            self.status.emit(text)
            #time.sleep(.1)
        self.update.emit(100)
        self.update.emit(0)
        self.status.emit("Installing boot loader...")

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

        print "Distro extraction 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(open(gen_fun.resource_path(os.path.join("tools", "version.txt")), 'r').read().strip())
if platform.system() == "Linux":
    window.setWindowIcon(QtGui.QIcon(gen_fun.resource_path(os.path.join("tools", "multibootusb.png"))))
else:
    window.setWindowIcon(QtGui.QIcon(gen_fun.resource_path(os.path.join("tools", "multibootusb.ico"))))
sys.exit(app.exec_())