#!/usr/bin/python3

from unidecode import unidecode
from subprocess import Popen, PIPE
import getopt
import gettext
import gi
import locale
import os
import signal
import subprocess
import sys

gi.require_version('Polkit', '1.0')
gi.require_version('Gtk', '3.0')
gi.require_version('UDisks', '2.0')
gi.require_version('XApp', '1.0')

from gi.repository import Polkit, Gtk, GLib, UDisks, XApp

try:
    gi.require_version('Unity', '7.0')
    from gi.repository import Unity
    Using_Unity = True
except Exception:
    Using_Unity = False

if Using_Unity:
    launcher = Unity.LauncherEntry.get_for_desktop_id("mintstick.desktop")

APP = 'mintstick'
LOCALE_DIR = "/usr/share/locale"
locale.bindtextdomain(APP, LOCALE_DIR)
gettext.bindtextdomain(APP, LOCALE_DIR)
gettext.textdomain(APP)
_ = gettext.gettext

# https://technet.microsoft.com/en-us/library/bb490925.aspx
FORBIDDEN_CHARS = ["*", "?", "/", "\\", "|", ".", ",", ";", ":", "+", "=", "[", "]", "<", ">", "\""]

RELEVANT_UDISK_PROPERTIES = ['connection-bus', 'ejectable', 'id', \
'media-available', 'media-compatibility', 'media-removable', \
'model', 'vendor', 'optical', 'removable', 'size']

# noinspection PyUnusedLocal
class MintStick:
    def __init__(self, iso_path_arg=None, usb_path_arg=None, filesystem_arg=None, mode_arg=None, debug_arg=False):

        self.debug = debug_arg
        self.filesystem = filesystem_arg

        def devices_changed_callback(client):
            self.get_devices()

        self.udisks_client = UDisks.Client.new_sync()
        self.udisk_listener_id = self.udisks_client.connect("changed", devices_changed_callback)

        self.wTree = Gtk.Builder()
        self.wTree.set_translation_domain(APP)
        self.wTree.add_from_file("/usr/share/mintstick/mintstick.ui")

        self.process = None
        self.source_id = None
        self.dev = None
        self.write_progress = None

        if mode_arg == "iso":
            self.mode = "normal"
            self.devicelist = self.wTree.get_object("device_combobox")
            self.go_button = self.wTree.get_object("write_button")
            self.verify_button = self.wTree.get_object("verify_button")
            self.progressbar = self.wTree.get_object("progressbar")
            self.chooser = self.wTree.get_object("filechooserbutton")

            # Making file chooser accessible for users with screen reader
            label = self.wTree.get_object("label_write_image")
            button = self.chooser.get_children()[0]
            label.set_mnemonic_widget(button)

            # Devicelist model
            self.devicemodel = Gtk.ListStore(str, str)

            # Renderer
            renderer_text = Gtk.CellRendererText()
            self.devicelist.pack_start(renderer_text, True)
            self.devicelist.add_attribute(renderer_text, "text", 1)

            self.get_devices()
            # get globally needed widgets
            self.window = self.wTree.get_object("main_window")
            self.window.connect("destroy", self.close)

            # set default file filter to *.iso/*.img
            filt = Gtk.FileFilter()
            filt.add_pattern("*.[iI][mM][gG]")
            filt.add_pattern("*.[iI][sS][oO]")
            self.chooser.set_filter(filt)

            self.devicelist.connect("changed", self.device_selected)
            self.go_button.connect("clicked", self.do_write)
            self.verify_button.connect("clicked", self.verify)
            self.chooser.connect("file-set", self.file_selected)

            if iso_path_arg:
                if os.path.exists(iso_path_arg):
                    self.chooser.set_filename(iso_path_arg)
                    self.file_selected(self.chooser)
                    self.verify_button.set_sensitive(True)

        if mode_arg == "format":
            self.mode = "format"
            self.devicelist = self.wTree.get_object("formatdevice_combobox")
            self.go_button = self.wTree.get_object("format_formatbutton")
            self.label_entry = self.wTree.get_object("volume_label_entry")
            try:
                self.label_entry.set_text(unidecode(self.label_entry.get_text()))
            except:
                self.label_entry.set_text("USB STICK")
            self.label_entry_changed_id = self.label_entry.connect("changed", self.on_label_entry_text_changed)

            self.window = self.wTree.get_object("format_window")
            self.window.connect("destroy", self.close)

            self.progressbar = self.wTree.get_object("format_progressbar")
            self.filesystemlist = self.wTree.get_object("filesystem_combobox")

            self.go_button.connect("clicked", self.do_format)
            self.filesystemlist.connect("changed", self.filesystem_selected)
            self.devicelist.connect("changed", self.device_selected)

            self.label_entry.connect("changed", lambda *_: self.update_format_button())
            self.filesystemlist.connect("changed", lambda *_: self.update_format_button())
            self.devicelist.connect("changed", lambda *_: self.update_format_button())

            # Filesystemlist
            self.fsmodel = Gtk.ListStore(str, str, int, bool, bool)
            #                     id       label    max-length force-upper-case   force-alpha-numeric
            self.fsmodel.append(["fat32", "FAT32", 11, True, True])
            self.fsmodel.append(["exfat", "exFAT", 15, False, False])
            self.fsmodel.append(["ntfs", "NTFS", 32, False, False])
            self.fsmodel.append(["ext4", "EXT4", 16, False, False])
            self.filesystemlist.set_model(self.fsmodel)

            # Renderer
            renderer_text = Gtk.CellRendererText()
            self.filesystemlist.pack_start(renderer_text, True)
            self.filesystemlist.add_attribute(renderer_text, "text", 1)

            # Devicelist model
            self.devicemodel = Gtk.ListStore(str, str)

            # Renderer
            renderer_text = Gtk.CellRendererText()
            self.devicelist.pack_start(renderer_text, True)
            self.devicelist.add_attribute(renderer_text, "text", 1)

            self.filesystemlist.set_sensitive(True)
            # Default's to fat32
            self.filesystemlist.set_active(0)
            if filesystem_arg is not None:
                fs_iter = self.fsmodel.get_iter_first()
                while fs_iter is not None:
                    value = self.fsmodel.get_value(fs_iter, 0)
                    if value == filesystem_arg:
                        self.filesystemlist.set_active_iter(fs_iter)
                    fs_iter = self.fsmodel.iter_next(fs_iter)

            self.get_devices()
            self.select_device(usb_path_arg)

        self.window.show()

    def select_device(self, usb_path=None):
        if usb_path is not None:
            device_iter = self.devicemodel.get_iter_first()
            while device_iter is not None:
                value = self.devicemodel.get_value(device_iter, 0)
                if usb_path in value:
                    self.devicelist.set_active_iter(device_iter)
                    return
                device_iter = self.devicemodel.iter_next(device_iter)
            return

        device_iter = self.devicemodel.get_iter_first()
        if device_iter is not None:
            self.devicelist.set_active_iter(device_iter)

    def verify(self, button):
        subprocess.Popen(["mint-iso-verify", self.chooser.get_filename()])

    def print_drive(self, drive):
        # print drive info to stdout
        # this is done for debugging purposes
        try:
            for prop in drive.list_properties():
                name = prop.name
                if name in RELEVANT_UDISK_PROPERTIES:
                    print(f"    {prop.name}: {drive.get_property(prop.name)}")
            print()
        except Exception as e:
            print(e)

    def get_devices(self):
        self.go_button.set_sensitive(False)
        self.devicemodel.clear()
        detected_drives = []
        self.dev = None

        manager = self.udisks_client.get_object_manager()

        for obj in manager.get_objects():
            if obj is not None:
                block = obj.get_block()
                if block is not None:
                    drive = self.udisks_client.get_drive_for_block(block)
                    if drive is not None:
                        drive_id = drive.get_property('id')
                        if drive_id in detected_drives:
                            continue
                        self.print_drive(drive)
                        is_usb = str(drive.get_property('connection-bus')) in ['usb', 'cpio', 'sdio']
                        size = int(drive.get_property('size'))
                        optical = bool(drive.get_property('optical'))
                        removable = bool(drive.get_property('removable'))

                        if is_usb and size > 0 and removable and not optical:
                            name = block.get_property('device')
                            name = ''.join([i for i in name if not i.isdigit()])

                            drive_vendor = str(drive.get_property('vendor'))
                            drive_model = str(drive.get_property('model'))

                            if drive_vendor.strip() != "":
                                drive_model = "%s %s" % (drive_vendor, drive_model)

                            if size >= 1000000000000:
                                size = "%.0fTB" % round(size / 1000000000000)
                            elif size >= 1000000000:
                                size = "%.0fGB" % round(size / 1000000000)
                            elif size >= 1000000:
                                size = "%.0fMB" % round(size / 1000000)
                            elif size >= 1000:
                                size = "%.0fkB" % round(size / 1000)
                            else:
                                size = "%.0fB" % round(size)

                            item = "%s (%s) - %s" % (drive_model, name, size)

                            detected_drives.append(drive_id)
                            self.devicemodel.append([name, item])

        self.devicelist.set_model(self.devicemodel)

    def device_selected(self, widget):
        iterator = self.devicelist.get_active_iter()
        if iterator is not None:
            self.dev = self.devicemodel.get_value(iterator, 0)
            self.go_button.set_sensitive(True)

    def filesystem_selected(self, widget):
        fs_iter = self.filesystemlist.get_active_iter()
        if fs_iter is not None:
            self.filesystem = self.fsmodel.get_value(fs_iter, 0)
            self.activate_devicelist()

            self.label_entry.set_max_length(self.fsmodel.get_value(fs_iter, 2))
            self.on_label_entry_text_changed(self, self.label_entry)

    def file_selected(self, widget):
        self.activate_devicelist()
        filename = self.chooser.get_filename()
        if filename is not None and os.path.exists(filename):
            self.verify_button.set_sensitive(True)
        else:
            self.verify_button.set_sensitive(False)

    def fix_label_entry_text(self, should_block_handler):
        if should_block_handler:
            self.label_entry.handler_block(self.label_entry_changed_id)

        active_iter = self.filesystemlist.get_active_iter()
        text = self.label_entry.get_text()

        if self.fsmodel.get_value(active_iter, 3):
            text = text.upper()

        if self.fsmodel.get_value(active_iter, 4):
            for char in FORBIDDEN_CHARS:
                text = text.replace(char, "")

        try:
            text = unidecode(text)
        except Exception as e:
            print(f"Exception in sanitizing the label {text}")
            print(str(e))
            text = "USB STICK"

        self.label_entry.set_text(text)

        if should_block_handler:
            self.label_entry.handler_unblock(self.label_entry_changed_id)

    def on_label_entry_text_changed(self, widget, data=None):
        self.fix_label_entry_text(True)

    def update_format_button(self):
        has_valid_device = self.devicelist.get_active_iter() is not None
        has_valid_fs = self.filesystemlist.get_active_iter() is not None
        is_enough_length = self.label_entry.get_buffer().get_length() > 0

        if has_valid_device and has_valid_fs and is_enough_length:
            self.go_button.set_sensitive(True)
        else:
            self.go_button.set_sensitive(False)


    def do_format(self, widget):
        if self.debug:
            print("DEBUG: Format %s as %s" % (self.dev, self.filesystem))
            return
        self.udisks_client.handler_block(self.udisk_listener_id)
        self.devicelist.set_sensitive(False)
        self.filesystemlist.set_sensitive(False)
        self.go_button.set_sensitive(False)
        self.label_entry.set_sensitive(False)
        label = self.label_entry.get_text()
        self.raw_format(self.dev, self.filesystem, label)


    def check_format_job(self):
        self.process.poll()
        if self.process.returncode is None:
            self.pulse_progress()
            return True
        else:
            GLib.idle_add(self.format_job_done, self.process.returncode)
            self.process = None
            return False

    def raw_format(self, usb_path_arg, fstype, label):
        if os.geteuid() > 0:
            polkit_exec = 'pkexec'
            self.process = Popen(
                [polkit_exec, '/usr/bin/mint-stick-format', '-d', usb_path_arg,
                 '-f', fstype, '-u', str(os.geteuid()), '-g', str(os.getgid()), '--', label],
                shell=False, stdout=PIPE, preexec_fn=os.setsid)
        else:
            self.process = Popen(
                ['/usr/bin/mint-stick-format', '-d', usb_path_arg,
                 '-f', fstype, '-u', str(os.geteuid()), '-g', str(os.getgid()), "--", label],
                shell=False, stdout=PIPE, preexec_fn=os.setsid)

        self.progressbar.show()
        self.pulse_progress()

        GLib.timeout_add(500, self.check_format_job)

    def format_job_done(self, rc):
        if rc == 0:
            self.show_format_result("dialog-information", _('The USB stick was formatted successfully.'))
            return False
        elif rc == 5:
            message = _("An error occured while creating a partition on %s.") % usb_path
        elif rc == 127:
            message = _('Authentication Error.')
        elif rc == 126:  # Cancelled by the user.
            self.clear_progress()
            message = None
        else:
            message = _('An error occurred.')

        if message is not None:
            self.show_format_result("dialog-error", message)
        self.set_format_sensitive(False)
        self.udisks_client.handler_unblock(self.udisk_listener_id)
        return False

    def do_write(self, widget):
        if self.debug:
            print("DEBUG: Write %s to %s" % (self.chooser.get_filename(), self.dev))
            return

        source = self.chooser.get_filename()
        target = self.dev
        filename = os.path.basename(source).lower()
        # Don't write Windows ISO
        for keyword in ["windows", "win7", "win8", "win10", "winxp"]:
            if keyword in filename:
                self.wTree.get_object("stack").set_visible_child_name("windows_page")
                return

        self.udisks_client.handler_block(self.udisk_listener_id)
        self.go_button.set_sensitive(False)
        self.verify_button.set_sensitive(False)
        self.devicelist.set_sensitive(False)
        self.chooser.set_sensitive(False)
        self.progressbar.show()
        self.raw_write(source, target)

    def set_progress(self, size):
        self.progressbar.set_fraction(size)
        str_progress = "%3.0f%%" % (float(size) * 100)
        int_progress = int(float(size) * 100)
        self.progressbar.set_text(str_progress)
        XApp.set_window_progress_pulse(self.window, False)
        XApp.set_window_progress(self.window, int_progress)

    def clear_progress(self):
        self.progressbar.hide()
        XApp.set_window_progress_pulse(self.window, False)

    def pulse_progress(self):
        self.progressbar.pulse()
        XApp.set_window_progress_pulse(self.window, True)

    def update_progress(self, fd, condition):
        if Using_Unity:
            launcher.set_property("progress_visible", True)
        if condition is GLib.IO_IN:
            line = fd.readline()
            # noinspection PyBroadException
            try:
                size = float(line.strip())
                progress = round(size * 100)
                if progress > self.write_progress:
                    self.write_progress = progress
                    GLib.idle_add(self.set_progress, size)
                    if Using_Unity:
                        launcher.set_property("progress", size)
            except:
                pass
            return True
        else:
            GLib.source_remove(self.source_id)
            return False

    def check_write_job(self):
        self.process.poll()
        if self.process.returncode is None:
            return True
        else:
            GLib.idle_add(self.write_job_done, self.process.returncode)
            self.process = None
            return False

    def raw_write(self, source, target):
        if os.geteuid() > 0:
            polkit_exec = 'pkexec'
            self.process = Popen(
                [polkit_exec, '/usr/bin/mint-stick-write', '-s', source, '-t', target],
                shell=False, stdout=PIPE, preexec_fn=os.setsid)
        else:
            self.process = Popen(
                ['/usr/bin/mint-stick-write', '-s', source, '-t', target],
                shell=False, stdout=PIPE, preexec_fn=os.setsid)

        self.write_progress = 0
        self.source_id = GLib.io_add_watch(self.process.stdout, GLib.IO_IN | GLib.IO_HUP, self.update_progress)
        GLib.timeout_add(500, self.check_write_job)

    def write_job_done(self, rc):
        self.udisks_client.handler_unblock(self.udisk_listener_id)
        if rc == 0:
            if Using_Unity:
                launcher.set_property("progress_visible", False)
                launcher.set_property("urgent", True)
            self.set_progress(1.0)
            self.show_result("dialog-information", _('The image was successfully written.'))
            return False
        elif rc == 3:
            message = _('Not enough space on the USB stick.')
        elif rc == 4:
            message = _('An error occured while copying the image.')
        elif rc == 127:
            message = _('Authentication Error.')
        elif rc == 126:  # Cancelled by the user.
            self.clear_progress()
            self.set_iso_sensitive()
            message = None
        else:
            message = _('An error occurred.')
        if message is not None:
            self.show_format_result("dialog-error", message)

        return False

    def show_format_result(self, icon_name, text):
        self.wTree.get_object("format_stack").set_visible_child_name("format_result_page")
        self.wTree.get_object("format_result_image").set_from_icon_name(icon_name, Gtk.IconSize.DIALOG)
        self.wTree.get_object("format_result_label").set_text(text)

    def show_result(self, icon_name, text):
        self.wTree.get_object("stack").set_visible_child_name("result_page")
        self.wTree.get_object("result_image").set_from_icon_name(icon_name, Gtk.IconSize.DIALOG)
        self.wTree.get_object("result_label").set_text(text)

    def close(self, widget):
        if self.process is not None:
            try:
                os.killpg(self.process.pid, signal.SIGTERM)
            except:
                pass
            finally:
                Gtk.main_quit()
        else:
            Gtk.main_quit()

    def activate_devicelist(self):
        self.devicelist.set_sensitive(True)

    def set_iso_sensitive(self):
        self.chooser.set_sensitive(True)
        self.devicelist.set_sensitive(True)
        self.go_button.set_sensitive(True)
        self.verify_button.set_sensitive(True)

    def set_format_sensitive(self, reset=True):
        if reset:
            self.get_devices()
        self.filesystemlist.set_sensitive(True)
        self.devicelist.set_sensitive(True)
        self.label_entry.set_sensitive(True)
        self.go_button.set_sensitive(True)

if __name__ == "__main__":
    usb_path = None
    iso_path = None
    filesystem = None
    mode = None

    def usage():
        print("Usage: mintstick [--debug] -m [format|iso]              : mode (format usb stick or burn iso image)")
        print("       mintstick [--debug] -m iso [-i|--iso] iso_path")
        print("       mintstick [--debug] -m format [-u|--usb] usb_device ")
        print("                           [-f|--filesystem] filesystem")
        exit(0)

    try:
        opts, args = getopt.getopt(sys.argv[1:], "hm:i:u:f:", ["debug", "help", "mode=", "iso=", "usb=", "filesystem="])
    except getopt.error as msg:
        print(msg)
        print("for help use --help")
        sys.exit(2)

    debug = False
    for o, a in opts:
        if o in ("-h", "--help"):
            usage()
        elif o in ("-i", "--iso"):
            iso_path = a
        elif o in ("-u", "--usb"):
            # hack. KDE Solid application gives partition name not device name.
            # Need to remove extra digit from the string.
            # ie. /dev/sdj1 -> /dev/sdj
            usb_path = ''.join([i for i in a if not i.isdigit()])
        elif o in ("-f", "--filesystem"):
            filesystem = a
        elif o in ("-m", "--mode"):
            mode = a
        elif o == "--debug":
            debug = True

    argc = len(sys.argv)
    if argc > 8:
        print("Too many arguments")
        print("for help use --help")
        exit(2)

    # Mandatory argument
    if (mode is None) or ((mode != "format") and (mode != "iso")):
        usage()

    MintStick(iso_path, usb_path, filesystem, mode, debug)
    Gtk.main()