#!/usr/bin/python3

import sys, os, apt
import subprocess
import filecmp

ORIGIN = "LMDE 3 'Cindy'"
ORIGIN_CODENAME = "cindy"
ORIGIN_BASE_CODENAME = "stretch"

DESTINATION = "LMDE 4 'Debbie'"
DESTINATION_CODENAME = "debbie"
DESTINATION_BASE_CODENAME = "buster"

SUPPORTED_EDITIONS = ["cinnamon"]

CHECK_ABSENT = ["libservlet3.1-java"]
CHECK_PRESENT = ["default-jre"]
CHECK_UP_TO_DATE = ["mintupgrade", "apt", "dpkg", "linuxmint-keyring", "debian-archive-keyring", "mintsystem"]

BACKUP_APT_SOURCES = os.path.expanduser("~/Upgrade-Backup/APT/")
BACKUP_FSTAB = os.path.expanduser("~/Upgrade-Backup/fstab")

PACKAGES_PRE_REMOVALS = []
PACKAGES_REMOVALS = ["tomboy", "libxplayer-plparser18", "xplayer-common", "gksu", "memtest86+"]
PACKAGES_ADDITIONS = ["neofetch", "ffmpegthumbnailer", "amd64-microcode", "intel-microcode", "celluloid", "drawing", "gnote", "grub2-theme-mint"]

IMPORTANT_PACKAGES = ["cinnamon", "xreader", "xed", "mintsystem", "metacity", "nemo", "nemo-preview"]

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

class MintUpgrade():

    def __init__(self):

        self.reversible = True

        # Check the Mint info file
        if not os.path.exists("/etc/linuxmint/info"):
            self.fail("Missing file '/etc/linuxmint/info'.")

        # Check the edition
        self.mint_codename = 'unknown'
        self.mint_edition = 'unknown'
        self.mint_meta = 'unknown'
        with open("/etc/linuxmint/info", "r") as info:
            for line in info:
                line = line.strip()
                if "EDITION=" in line:
                    self.mint_edition = line.split('=')[1].replace('"', '').split()[0]
                    self.mint_meta = "mint-meta-%s" % self.mint_edition.lower()
                if "CODENAME=" in line:
                    self.mint_codename = line.split('=')[1].replace('"', '').split()[0]
        self.points_to_destination = False
        if os.path.exists("/etc/apt/sources.list.d/official-package-repositories.list"):
            with open("/etc/apt/sources.list.d/official-package-repositories.list") as sources:
                for line in sources:
                    if DESTINATION_CODENAME in line:
                        self.points_to_destination = True
                        break

    def restore_sources(self):
        self.reversible = False
        self.progress("Restoring your backed up APT sources")
        if os.path.exists(BACKUP_APT_SOURCES):
            self.check_command("sudo mkdir -p /etc/apt/sources.list.d", "Failed to restore APT sources")
            self.check_command("sudo rm -rf /etc/apt/sources.list.d/*", "Failed to restore APT sources")
            self.check_command("sudo cp -R %s/* /etc/apt/" % BACKUP_APT_SOURCES, "Failed to restore APT sources")
            self.check_command("sudo rm -rf '%s'" % BACKUP_APT_SOURCES, "Failed to restore APT sources")

    def prepare(self):
        # Check codename
        self.progress("Checking your Linux Mint codename")
        if self.mint_codename != ORIGIN_CODENAME and self.mint_codename != DESTINATION_CODENAME:
            self.fail("Your version of Linux Mint is '%s'. Only %s can be upgraded to %s." % (self.mint_codename.capitalize(), ORIGIN, DESTINATION))

        # Check edition
        self.progress("Checking your Linux Mint edition")
        if self.mint_edition.lower() not in SUPPORTED_EDITIONS:
            self.fail("Your edition of Linux Mint is '%s'. It cannot be upgraded to %s." % (self.mint_edition, DESTINATION))

        # Check for timeshift configuration
        self.progress("Checking your Timeshift configuration")
        if not os.path.exists("/etc/timeshift.json"):
            self.fail("Please set up system snapshots. If anything goes wrong with the upgrade, snapshots will allow you to restore your operating system. Install and configure Timeshift, and create a snapshot before proceeding with the upgrade.")

        if not self.points_to_destination:
            # Check packages
            self.progress("Checking packages")
            cache = apt.Cache()
            need_to_remove = []
            for pkg in CHECK_ABSENT:
                if pkg in cache:
                    pkg = cache[pkg]
                    if pkg.is_installed:
                        need_to_remove.append(pkg.name)
            if len(need_to_remove) > 0:
                self.fail("The following packages create issues with this upgrade, please remove them: %s" % ", ".join(need_to_remove))
            need_to_install = []
            for pkg in CHECK_PRESENT:
                if pkg in cache:
                    pkg = cache[pkg]
                    if not pkg.is_installed:
                        need_to_install.append(pkg.name)
            if len(need_to_remove) > 0:
                self.fail("The following packages are required for a smooth install, please install them: %s" % ", ".join(need_to_install))

            self.progress("Updating cache")
            os.system("DEBIAN_PRIORITY=critical sudo apt-get update")
            cache = apt.Cache()

            # Check that we don't have any Debian Multimedia packages
            self.progress("Checking deb-multimedia.org repository")
            output = subprocess.check_output("apt policy", shell=True).decode("UTF-8")
            if "deb-multimedia.org" in output:
                self.fail("The Unofficial Debian Multimedia repository (https://deb-multimedia.org) is no longer used in LMDE 4. Please remove it from your APT configuration (usually in /etc/apt/sources.list.d/official-package-repositories.list).")

            self.progress("Checking deb-multimedia.org packages")
            dmo_packages = []
            for pkg in cache:
                if pkg.is_installed:
                    if "-dmo" in pkg.installed.version:
                        dmo_packages.append(pkg.name)
                    if pkg.installed.source_name.endswith("-dmo"):
                        dmo_packages.append(pkg.name)
            if len(dmo_packages) > 0:
                self.fail("The following packages come from the Unofficial Debian Multimedia repository: \n\n     %s.\n\nLaunch mintsources and go to the \"Maintenance tab\".\nFirst, use \"Downgrade Foreign Packages\" to downgrade all the packages.\nThen, use \"Remove Foreign Packages\" to remove any remaining packages with '-dmo' in their version number.\n" % ", ".join(dmo_packages))

            # Check that we're up to date
            self.progress("Checking if Linux Mint is up to date")
            for pkg in CHECK_UP_TO_DATE:
                if pkg in cache:
                    pkg = cache[pkg]
                    if pkg.is_installed and pkg.installed.version != pkg.candidate.version:
                        self.fail("Your operating system is not up to date. Please apply available updates and reboot the computer.")

        # Switch to the destination APT sources
        if not os.path.exists(BACKUP_APT_SOURCES):
            self.progress("Backing up your APT sources")
            messages = []
            messages.append("Your repositories will now be switched to point to %s." % DESTINATION)
            messages.append("Any 3rd party repositories or PPA will be removed.")
            messages.append("A backup of your APT sources will be written to %s." % BACKUP_APT_SOURCES)
            self.continue_yes_no(messages)
            os.system("mkdir -p %s" % BACKUP_APT_SOURCES)
            os.system("cp -R /etc/apt/sources.* %s/" % BACKUP_APT_SOURCES)
        self.progress("Setting up the repositories for %s" % DESTINATION)
        if os.path.exists("/etc/apt/sources.list"):
            self.check_command("sudo truncate --size 0 /etc/apt/sources.list", "Failed to configure APT sources")
        self.check_command("sudo mkdir -p /etc/apt/sources.list.d", "Failed to configure APT sources")
        self.check_command("sudo rm -rf /etc/apt/sources.list.d/*", "Failed to configure APT sources")
        self.check_command("sudo cp /usr/share/linuxmint/mintupgrade/apt_destination_sources /etc/apt/sources.list.d/official-package-repositories.list", "Failed to configure APT sources")
        self.try_command(2, 'DEBIAN_PRIORITY=critical sudo apt-get update', [])

    def check(self):
        self.progress("Simulating an upgrade")
        messages = []
        messages.append("APT will now calculate the package changes necessary to upgrade to %s." % DESTINATION)
        messages.append("If conflicts are detected and APT is unable to perform the upgrade, take note of the packages causing the issue, remove them, and re-install them after the upgrade.")
        messages.append("Pay close attention to what appears on the screen, and review the list of packages being REMOVED during the upgrade.")
        messages.append("Take note of the packages being removed, so you can eventually reinstall them after the upgrade.")
        self.continue_yes_no(messages)
        os.system('DEBIAN_PRIORITY=critical sudo apt-get dist-upgrade -o Dpkg::Options::="--force-confnew" -o Dpkg::Options::="--force-overwrite" --assume-no')

        cache = apt.Cache()
        cache.upgrade(True)
        changes = cache.get_changes()
        incorrect_removals = []
        kept_packages = []
        for pkg in changes:
            if pkg.is_installed:
                if pkg.marked_keep:
                    kept_packages.append(pkg.name)
                elif pkg.marked_delete and pkg.name in IMPORTANT_PACKAGES:
                    incorrect_removals.append(pkg.name)
        if len(incorrect_removals) > 0:
            self.restore_sources()
            self.fail("Performing the upgrade would remove the following important packages: %s." % ", ".join(sorted(incorrect_removals)))
        if len(kept_packages) > 0:
            self.warn("The following packages will be kept back during the upgrade: %s. \n\nThis might or might not indicate a problem. Check the APT output above to decide whether to continue with the upgrade. If this OK you might need to update them manually after the upgrade." % ", ".join(sorted(kept_packages)))

    def download(self):
        self.progress("Downloading upgrade packages")
        messages = []
        messages.append("APT will now download the package updates necessary for the upgrade to %s." % DESTINATION)
        self.continue_yes_no(messages)
        self.check_command("DEBIAN_PRIORITY=critical sudo apt-get dist-upgrade --download-only --yes", "Failed to download packages for the upgrade.")

    def upgrade(self):
        if self.mint_edition.lower() == "cinnamon":
            self.progress("Disabling the Cinnamon screensaver")
            os.system("killall cinnamon-screensaver")
        elif self.mint_edition.lower() == "mate":
            self.progress("Disabling the MATE screensaver")
            os.system("killall mate-screensaver")

        self.progress("Saving /etc/fstab")
        os.system("cp /etc/fstab %s" % BACKUP_FSTAB)

        self.progress("Removing blacklisted packages")
        for removal in PACKAGES_PRE_REMOVALS:
            os.system('sudo apt-get remove --yes %s' % removal) # The return code indicates a failure if some packages were not found, so ignore it.

        self.progress("Performing upgrade")
        messages = []
        messages.append("APT will perform the upgrade to %s." % DESTINATION)
        messages.append("This operation is non-reversible.")
        messages.append("Make sure you made backups, you tested %s in live mode and you performed your favorite superstitious tricks before proceeding." % DESTINATION)
        self.continue_yes_no(messages)

        self.reversible = False

        # Disable mintsystem during the upgrade
        os.system("sudo crudini --set /etc/linuxmint/mintSystem.conf global enabled False")

        fallback_commands = []
        fallback_commands.append("sudo dpkg --configure -a")
        fallback_commands.append("sudo apt-get install -fyq")

        result = self.try_command(5, 'DEBIAN_FRONTEND=noninteractive DEBIAN_PRIORITY=critical sudo apt-get upgrade -fyq -o Dpkg::Options::="--force-confnew" -o Dpkg::Options::="--force-overwrite"', fallback_commands)
        if not result:
            self.progress("An issue was detected during the upgrade, running the upgrade in manual mode.")
            self.check_command('sudo apt-get upgrade -o Dpkg::Options::="--force-confnew" -o Dpkg::Options::="--force-overwrite"', "Failed to upgrade some of the packages. Please review the error message, use APT to fix the situation and try again.")

        result = self.try_command(5, 'DEBIAN_FRONTEND=noninteractive DEBIAN_PRIORITY=critical sudo apt-get dist-upgrade -fyq -o Dpkg::Options::="--force-confnew" -o Dpkg::Options::="--force-overwrite"', fallback_commands)
        if not result:
            self.progress("An issue was detected during the upgrade, running dist-upgrade in manual mode.")
            self.check_command('sudo apt-get dist-upgrade -o Dpkg::Options::="--force-confnew" -o Dpkg::Options::="--force-overwrite"', "Failed to dist-upgrade some of the packages. Please review the error message, use APT to fix the situation and try again.")

        self.progress("Re-installing the meta-package for your edition of Linux Mint")
        self.check_command('sudo apt-get install --yes %s' % self.mint_meta, "Failed to install %s" % self.mint_meta)

        # Enable APT recommends
        self.progress("Re-enabling APT recommends")
        if os.path.exists("/etc/apt/apt.conf.d/00recommends"):
            os.system("sudo rm -f /etc/apt/apt.conf.d/00recommends")

        self.progress("Re-installing the multimedia codecs")
        self.check_command('sudo apt-get install --yes mint-meta-codecs', "Failed to install mint-meta-codecs")

        self.progress("Installing new packages")
        self.check_command('sudo apt-get install --yes %s' % " ".join(PACKAGES_ADDITIONS), "Failed to install additional packages.")

        self.progress("Removing obsolete packages")
        for removal in PACKAGES_REMOVALS:
            os.system('sudo apt-get purge --yes %s' % removal) # The return code indicates a failure if some packages were not found, so ignore it.

        self.progress("Performing system adjustments")
        os.system("sudo rm -f /etc/systemd/logind.conf")
        os.system("apt install --reinstall -o Dpkg::Options::=\"--force-confmiss\" systemd")
        os.system("sudo rm -f /etc/polkit-1/localauthority/50-local.d/com.ubuntu.enable-hibernate.pkla")

        # Re-enable the plymouth theme (specific to LMDE)
        os.system("sudo plymouth-set-default-theme mint-logo")
        os.system("sudo update-initramfs -u")

        if os.path.exists("/usr/share/ubuntu-system-adjustments/systemd/adjust-grub-title"):
            os.system("sudo /usr/share/ubuntu-system-adjustments/systemd/adjust-grub-title")
        elif os.path.exists("/usr/share/debian-system-adjustments/systemd/adjust-grub-title"):
            os.system("sudo /usr/share/debian-system-adjustments/systemd/adjust-grub-title")

        # Re-enable mintsystem
        os.system("sudo crudini --set /etc/linuxmint/mintSystem.conf global enabled True")
        os.system("sudo /usr/lib/linuxmint/mintsystem/mint-adjust.py")

        # Restore /etc/fstab if it was changed
        if not filecmp.cmp('/etc/fstab', BACKUP_FSTAB):
            os.system("cp /etc/fstab %s.upgraded" % BACKUP_FSTAB)
            os.system("sudo cp %s /etc/fstab" % BACKUP_FSTAB)
            self.warn("A package modified /etc/fstab during the upgrade. To ensure a successful boot, the upgrader restored your original /etc/fstab and saved the modified file in %d.upgraded." % BACKUP_FSTAB)

        self.progress("The upgrade is finished. Reboot the computer with \"sudo reboot\" when ready.")

    def check_command(self, command, message):
        ret = os.system(command)
        if ret != 0:
            self.fail(message)

    def try_command(self, num_times, command, fallback_commands):
        success = False
        for i in range(num_times):
            ret = os.system(command)
            if ret == 0:
                return True
            self.progress("Error detected on try #%d..." % (i+1))
            if (i+1) < num_times:
                self.progress("Retrying...")
            for fallback_command in fallback_commands:
                self.progress ("Running fallback command '%s'" % fallback_command)
                os.system(fallback_command)

    def fail(self, message):
        print ("")
        print ("------------------------------------------------")
        print ("%s!!  ERROR: %s%s" % (bcolors.FAIL, message, bcolors.ENDC))
        print ("!!  Exiting.")
        print ("------------------------------------------------")
        print ("")
        if self.reversible:
            self.restore_sources()
        sys.exit(1)

    def continue_yes_no(self, messages):
        print ("")
        print ("-------------------------------------------------")
        for message in messages:
            print ("%s    %s%s" % (bcolors.WARNING, message, bcolors.ENDC))
        print ("")
        answer = None
        while (answer not in ["y", "yes", "n", "no"]):
            answer = input("%s    Do you want to continue? [y/n]:%s " % (bcolors.OKGREEN, bcolors.ENDC)).lower()
        if answer in ["n", "no"]:
            print ("Exiting.")
            if self.reversible:
                self.restore_sources()
            sys.exit(0)

    def progress(self, message):
        print ("")
        print ("%s  + %s...%s" % (bcolors.HEADER, message, bcolors.ENDC))

    def warn(self, message):
        print ("")
        print ("%s  + %s%s" % (bcolors.WARNING, message, bcolors.ENDC))

def usage():
    print ("")
    print ("%sUsage:%s mintupgrade command" % (bcolors.HEADER, bcolors.ENDC))
    print ("")
    print ("%sCommands:%s" % (bcolors.HEADER, bcolors.ENDC))
    print ("  help                   - prints this usage note")
    print ("  check                  - checks the upgrade to %s" % DESTINATION)
    print ("  prepare                - prepares the upgrade to %s" % DESTINATION)
    print ("  download               - downloads the packages for the upgrade to %s" % DESTINATION)
    print ("  upgrade                - upgrades to %s" % DESTINATION)
    print ("  restore-sources        - restores the backed up APT sources (only use this command if you're still running %s)" % ORIGIN)
    print ("")
    sys.exit(0)

if __name__ == '__main__':

    if os.getuid() == 0:
        print ("")
        print ("Please don't run this command as root or with elevated privileges.")
        print ("")
        sys.exit(1)

    os.system("clear")

    if len(sys.argv) != 2:
        usage()
    command = sys.argv[1]
    if command == "help":
        usage()

    upgrader = MintUpgrade()

    if command == "restore-sources":
        upgrader.restore_sources()
    elif command == "check":
        upgrader.prepare()
        upgrader.check()
        upgrader.restore_sources()
    elif command == "prepare":
        upgrader.prepare()
    elif command == "download":
        upgrader.prepare()
        upgrader.download()
    elif command == "upgrade":
        upgrader.prepare()
        upgrader.download()
        upgrader.upgrade()
    else:
        usage()
