#!/usr/bin/env python
""" Helper script for the Packages plugin, used if yum library support
is enabled.  The yum libs have horrific memory leaks, so apparently
the right way to get around that in long-running processes it to have
a short-lived helper.  No, seriously -- check out the yum-updatesd
code.  It's pure madness. """

__revision__ = '$Revision$'

import os
import sys
import yum
import logging
import Bcfg2.Logger
from optparse import OptionParser, OptionError

try:
    import json
except ImportError:
    import simplejson as json

logger = logging.getLogger('bcfg2-yum-helper')

class DepSolver(object):
    def __init__(self, cfgfile, verbose=1):
        self.cfgfile = cfgfile
        self.yumbase = yum.YumBase()
        try:
            self.yumbase.preconf.debuglevel = verbose
            self.yumbase.preconf.fn = cfgfile
        except AttributeError:
            self.yumbase._getConfig(cfgfile, debuglevel=verbose)

    def get_groups(self):
        try:
            return self._groups
        except AttributeError:
            return ["noarch"]

    def set_groups(self, groups):
        self._groups = set(groups).add("noarch")

    groups = property(get_groups, set_groups)

    def is_package(self, package):
        if isinstance(package, tuple):
            if package[1] is None and package[2] == (None, None, None):
                package = package[0]
            else:
                return None

        return bool(self.get_package_object(package, silent=True))

    def is_virtual_package(self, package):
        return bool(self.get_provides(package, silent=True))

    def get_package_object(self, package, silent=False):
        try:
            matches = self.yumbase.pkgSack.returnNewestByName(name=package)
        except yum.Errors.PackageSackError:
            if not silent:
                logger.warning("Packages: Package '%s' not found" %
                               self.get_package_name(package))
            matches = []
        except yum.Errors.RepoError:
            err = sys.exc_info()[1]
            logger.error("Packages: Temporary failure loading metadata for "
                         "'%s': %s" % (self.get_package_name(package), err))
            matches = []

        pkgs = self._filter_arch(matches)
        if pkgs:
            return pkgs[0]
        else:
            return None

    def get_deps(self, package):
        pkg = self.get_package_object(package)
        deps = []
        if pkg:
            deps = set(pkg.requires)
            # filter out things the package itself provides
            deps.difference_update([dep for dep in deps
                                        if pkg.checkPrco('provides', dep)])
        else:
            logger.error("Packages: No package available: %s" %
                         self.get_package_name(package))
        return deps

    def get_provides(self, required, all=False, silent=False):
        if not isinstance(required, tuple):
            required = (required, None, (None, None, None))

        try:
            prov = \
                self.yumbase.whatProvides(*required).returnNewestByNameArch()
        except yum.Errors.NoMoreMirrorsRepoError:
            err = sys.exc_info()[1]
            logger.error("Packages: Temporary failure loading metadata for "
                         "'%s': %s" %
                         (self.get_package_name(required), err))
            return []

        if prov and not all:
            prov = self._filter_provides(required, prov)
        elif not prov and not silent:
            logger.error("Packages: No package provides %s" %
                         self.get_package_name(required))
        return prov

    def get_group(self, group, ptype="default"):
        if group.startswith("@"):
            group = group[1:]

        try:
            if self.yumbase.comps.has_group(group):
                group = self.yumbase.comps.return_group(group)
            else:
                logger.warning("Packages: '%s' is not a valid group" % group)
                return []
        except yum.Errors.GroupsError:
            err = sys.exc_info()[1]
            logger.warning("Packages: %s" % err)
            return []
        
        if ptype == "default":
            return [p
                    for p, d in list(group.default_packages.items())
                    if d]
        elif ptype == "mandatory":
            return [p
                    for p, m in list(group.mandatory_packages.items())
                    if m]
        elif ptype == "optional" or ptype == "all":
            return group.packages
        else:
            logger.warning("Unknown group package type '%s'" % ptype)
            return []

    def _filter_provides(self, package, providers):
        providers = [pkg for pkg in self._filter_arch(providers)]
        if len(providers) > 1:
            # go through each provider and make sure it's the newest
            # package of its name available.  If we have multiple
            # providers, avoid installing old packages.
            #
            # For instance: on Fedora 14,
            # perl-Sub-WrapPackages-2.0-2.fc14 erroneously provided
            # perl(lib), which should not have been provided;
            # perl(lib) is provided by the "perl" package.  The bogus
            # provide was removed in perl-Sub-WrapPackages-2.0-4.fc14,
            # but if we just queried to resolve the "perl(lib)"
            # dependency, we'd get both packages.  By performing this
            # check, we learn that there's a newer
            # perl-Sub-WrapPackages available, so it can't be the best
            # provider of perl(lib).
            rv = []
            for pkg in providers:
                if self.get_package_object(pkg.name) == pkg:
                    rv.append(pkg)
        else:
            rv = providers
        return [p.name for p in rv]

    def _filter_arch(self, packages):
        matching = [pkg for pkg in packages if pkg.arch in self.groups]
        if matching:
            return matching
        else:
            # no packages match architecture; we'll assume that the
            # user knows what s/he is doing and this is a multiarch
            # box.
            return packages

    def get_package_name(self, package):
        """ get the name of a package or virtual package from the
        internal representation used by this Collection class """
        if isinstance(package, tuple):
            return yum.misc.prco_tuple_to_string(package)
        else:
            return str(package)

    def complete(self, packagelist, groups=None):
        if groups is None:
            groups = []

        packages = set()
        pkgs = set(packagelist)
        requires = set()
        satisfied = set()
        unknown = set()
        final_pass = False

        while requires or pkgs:
            # infinite loop protection
            start_reqs = len(requires)
            
            while pkgs:
                package = pkgs.pop()
                if package in packages:
                    continue
                
                if not self.is_package(package):
                    # try this package out as a requirement
                    requires.add((package, None, (None, None, None)))
                    continue

                packages.add(package)
                reqs = set(self.get_deps(package)).difference(satisfied)
                if reqs:
                    requires.update(reqs)

            reqs_satisfied = set()
            for req in requires:
                if req in satisfied:
                    reqs_satisfied.add(req)
                    continue

                if req[1] is None and self.is_package(req[0]):
                    if req[0] not in packages:
                        pkgs.add(req[0])
                    reqs_satisfied.add(req)
                    continue
                    
                logger.debug("Packages: Handling requirement '%s'" %
                             self.get_package_name(req))
                providers = list(set(self.get_provides(req)))
                if len(providers) > 1:
                    # hopefully one of the providing packages is already
                    # included
                    best = [p for p in providers if p in packages]
                    if best:
                        providers = best
                    else:
                        # pick a provider whose name matches the requirement
                        best = [p for p in providers if p == req[0]]
                        if len(best) == 1:
                            providers = best
                        elif not final_pass:
                            # found no "best" package, so defer
                            providers = None
                        # else: found no "best" package, but it's the
                        # final pass, so include them all
                
                if providers:
                    logger.debug("Packages: Requirement '%s' satisfied by %s" %
                                 (self.get_package_name(req),
                                  ",".join([self.get_package_name(p)
                                            for p in providers])))
                    newpkgs = set(providers).difference(packages)
                    if newpkgs:
                        for package in newpkgs:
                            if self.is_package(package):
                                pkgs.add(package)
                            else:
                                unknown.add(package)
                    reqs_satisfied.add(req)
                elif providers is not None:
                    # nothing provided this requirement at all
                    unknown.add(req)
                    reqs_satisfied.add(req)
                # else, defer
            requires.difference_update(reqs_satisfied)

            # infinite loop protection
            if len(requires) == start_reqs and len(pkgs) == 0:
                final_pass = True

            if final_pass and requires:
                unknown.update(requires)
                requires = set()

        unknown = [self.get_package_name(p) for p in unknown]

        return packages, unknown

    def clean_cache(self):
        for mdtype in ["Headers", "Packages", "Sqlite", "Metadata",
                       "ExpireCache"]:
            # for reasons that are entirely obvious, all of the yum
            # API clean* methods return a tuple of 0 (zero, always
            # zero) and a list containing a single message about how
            # many files were deleted.  so useful.  thanks, yum.
            msg = getattr(self.yumbase, "clean%s" % mdtype)()[1][0]
            if not msg.startswith("0 "):
                logger.info("Packages: %s" % msg)


def main():
    parser = OptionParser()
    parser.add_option("-c", "--config", help="Config file")
    parser.add_option("-v", "--verbose", help="Verbosity level", action="count")
    (options, args) = parser.parse_args()
    try:
        cmd = args[0]
    except IndexError:
        logger.error("bcfg2-yum-helper: No command given")
        return 1

    if not os.path.exists(options.config):
        logger.error("bcfg2-yum-helper: Config file %s not found" %
                     options.config)
        return 1

    depsolver = DepSolver(options.config, options.verbose)
    if cmd == "clean":
        depsolver.clean_cache()
        print json.dumps(True)
    elif cmd == "complete":
        data = json.loads(sys.stdin.read())
        (packages, unknown) = depsolver.complete(data['packages'],
                                                 groups=data['groups'])
        print json.dumps(dict(packages=list(packages),
                              unknown=list(unknown)))
    elif cmd == "is_virtual_package":
        package = json.loads(sys.stdin.read())
        print json.dumps(bool(depsolver.get_provides(package, silent=True)))
    elif cmd == "get_deps" or cmd == "get_provides":
        package = json.loads(sys.stdin.read())
        print json.dumps(list(getattr(depsolver, cmd)(package)))
    elif cmd == "get_group":
        data = json.loads(sys.stdin.read())
        if "type" in data:
            packages = depsolver.get_group(data['group'], ptype=data['type'])
        else:
            packages = depsolver.get_group(data['group'])
        print json.dumps(list(packages))
    elif cmd == "is_package":
        package = json.loads(sys.stdin.read())
        print json.dumps(getattr(depsolver, cmd)(package))
        
        
if __name__ == '__main__':
    sys.exit(main())
