#!/usr/bin/python3
# -*- mode: python; coding: utf-8 -*-
#
# Copyright 2019, Cendio AB.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Library General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import sys
import pickle

try:
    import dnf
    import libdnf
except ImportError:
    sys.exit("Could not import DNF API")

def syntax():
    print("Syntax:", file=sys.stderr)
    print("    %s <command> [args]" % sys.argv[0], file=sys.stderr)
    sys.exit(1)

def do_test():
    # Don't need to do anything
    pass

def _prepare_transaction(base, pkgs):
    base.conf.substitutions.update_from_etc(base.conf.installroot)
    # Substitutions are needed for correct interpretation of repo files.
    RELEASEVER = dnf.rpm.detect_releasever(base.conf.installroot)
    base.conf.substitutions['releasever'] = RELEASEVER
    # Repositories are needed if we want to install anything.
    base.read_all_repos()
    # A sack is required by marking methods and dependency resolving.
    base.fill_sack()
    # Feature marking methods set the user request.
    for pkg in pkgs:
        try:
            base.install(pkg)
        except dnf.exceptions.MarkingError:
            sys.exit('No package found providing %s' % pkg)
    # Resolving finds a transaction that allows the packages installation.
    try:
        base.resolve()
    except dnf.exceptions.DepsolveError:
        sys.exit('Dependencies cannot be resolved.')

def _convert_pkgs(pkgs):
    output = []
    for pkg in pkgs:
        output.append({"name": pkg.name,
                       "version": "%s-%s" % (pkg.version, pkg.release),
                       "arch": pkg.arch})
    return output

def do_resolve(pkgs):
    with dnf.Base() as base:
        _prepare_transaction(base, pkgs)

        # FIXME: The documented functions installs() and removes() do
        #        not actually exist:
        #        https://bugzilla.redhat.com/show_bug.cgi?id=1709352

        # FIXME: This is what dnf itself does, but it is not part of
        #        the API documented to be stable:
        #        https://bugzilla.redhat.com/show_bug.cgi?id=1709359

        # Figure out what we asked for, what are dependencies, and what
        # needs to be removed
        resolved = set()
        depends = set()
        removed = set()
        for tsi in base.transaction:
            if tsi.action == libdnf.transaction.TransactionItemAction_REMOVE:
                removed.add(tsi)
            else:
                if tsi.reason == libdnf.transaction.TransactionItemReason_USER:
                    resolved.add(tsi)
                else:
                    depends.add(tsi)

        resolved = _convert_pkgs(resolved)
        depends = _convert_pkgs(depends)
        removed = _convert_pkgs(removed)

        data = { "resolved": resolved,
                 "depends": depends,
                 "removed": removed }
        # We need to use version 2 so that Python 2.x can read this
        sys.stdout.buffer.write(pickle.dumps(data, protocol=2))

class DownloadReporter(dnf.callback.DownloadProgress):
    def start(self, total_files, total_size, total_drpms=0):
        self.__total_size = total_size
        self.__payloads = {}

    def progress(self, payload, done):
        if payload not in self.__payloads:
            self.__payloads[payload] = [ 0, payload.download_size ]

        self.__payloads[payload][0] = done

        total_done = 0
        for key in self.__payloads:
            total_done += self.__payloads[key][0]

        percentage = total_done / self.__total_size
        print("DOWNLOAD: %d" % (100 * percentage),
               file=sys.stderr, flush=True)

class TransactionReporter(dnf.callback.TransactionProgress):
    def error(self, message):
        print("Error: %s" % message, file=sys.stderr, flush=True)

    def progress(self, package, action,
                 ti_done, ti_total, ts_done, ts_total):
        if action in [ dnf.callback.PKG_CLEANUP,
                       dnf.callback.PKG_DOWNGRADED,
                       dnf.callback.PKG_REINSTALLED,
                       dnf.callback.PKG_UPGRADED ]:
            message = "Cleanup %s..." % package
        elif action == dnf.callback.PKG_DOWNGRADE:
            message = "Downgrading %s..." % package
        elif action == dnf.callback.PKG_INSTALL:
            message = "Installing %s..." % package
        elif action in [ dnf.callback.PKG_OBSOLETE,
                         dnf.callback.PKG_OBSOLETED ]:
            message = "Obsoleting %s..." % package
        elif action == dnf.callback.PKG_REINSTALL:
            message = "Reinstalling %s..." % package
        elif action == dnf.callback.PKG_REMOVE:
            message = "Erasing %s..." % package
        elif action == dnf.callback.PKG_UPGRADE:
            message = "Upgrading %s..." % package
        elif action == dnf.callback.PKG_VERIFY:
            message = "Verifying %s..." % package
        elif action == dnf.callback.PKG_SCRIPTLET:
            message = "Running scriptlet for %s..." % package
        elif action == dnf.callback.TRANS_PREPARATION:
            message = "Preparing transaction..."
        elif action == dnf.callback.TRANS_POST:
            message = "Finishing transaction..."
        else:
            print("Warning: Unknown transaction action %s" % action,
                  file=sys.stderr, flush=True)
            return

        # The transaction actions are special and does not give us
        # meaningful progress information
        if action == dnf.callback.TRANS_PREPARATION:
            total_percentage = 0
            item_percentage = 0
        elif action == dnf.callback.TRANS_POST:
            total_percentage = 100
            item_percentage = 100
        elif action == dnf.callback.PKG_VERIFY:
            # FIXME: The verify steps are not counted in the total
            #        transaction steps, so we need to do some
            #        calculations to get this right. For now just
            #        assume it is negligable time...
            total_percentage = 100
            item_percentage = 100
        else:
            total_percentage = 100 * ts_done / ts_total
            item_percentage = 100 * ti_done / ti_total

        print("TRANSACTION: %d %d %s" % (total_percentage,
                                         item_percentage,
                                         message),
              file=sys.stderr, flush=True)

def do_install(pkgs):
    with dnf.Base() as base:
        _prepare_transaction(base, pkgs)

        # The packages to be installed must be downloaded first.
        try:
            base.download_packages(base.transaction.install_set,
                                   progress=DownloadReporter())
        except dnf.exceptions.DownloadError:
            sys.exit('Required package cannot be downloaded.')

        # The request can finally be fulfilled.
        base.do_transaction(display=TransactionReporter())

def main():
    if len(sys.argv) < 2:
        syntax()

    try:
        if sys.argv[1] == "test":
            do_test()
        elif sys.argv[1] == "resolve":
            do_resolve(sys.argv[2:])
        elif sys.argv[1] == "install":
            do_install(sys.argv[2:])
        else:
            sys.exit("Unknown command \"%s\"" % sys.argv[1])
    except dnf.exceptions.Error as e:
        sys.exit('DNF error: %s' % e)

if __name__ == "__main__":
    main()
