#!/usr/bin/python
#
# Output the package set that would be installed if the packages on the cmdline were installed
#

import os
import sys
import yum
from yum.Errors import YumBaseError
import argparse
import tempfile
import ConfigParser
import logging

def setup_argparse():
    parser = argparse.ArgumentParser(description="Output packages DNF has selected")

    # required arguments for image creation
    required = parser.add_argument_group("required arguments")
    required.add_argument("-r", "--release", help="release information", required=True, metavar="STRING")
    required.add_argument("-s", "--source", help="source repository (may be listed multiple times)",
                        metavar="REPOSITORY", action="append", default=[], required=True)
    parser.add_argument("-t", dest="pkg_file", help="Read packages from a Mako template file")
    parser.add_argument("-f", "--find", help="Find the package that pulls in this other package")
    parser.add_argument("--tempdir", help="Directory to store temporary DNF files")
    parser.add_argument("--proxy", help="Proxy URL to use for DNF")
    parser.add_argument("packages", help="Packages to Install", metavar="STRING", nargs='*', default=[])

    return parser


def get_ybo(tempdir, repositories, releasever, proxy=None):
    """ Create a yum Base object and setup the repositories and installroot

        :param list repositories: List of repositories to use for the installation
        :param string releasever: Release version to pass to dnf

    """
    def sanitize_repo(repo):
        """Convert bare paths to file:/// URIs, and silently reject protocols unhandled by yum"""
        if repo.startswith("/"):
            return "file://{0}".format(repo)
        elif any(repo.startswith(p) for p in ('http://', 'https://', 'ftp://', 'file://')):
            return repo
        else:
            return None

    # sanitize the repository urls, split out .repo files
    repo_urls = filter(None, [sanitize_repo(r) for r in repositories if not r.endswith(".repo")])
    repo_files = list(r for r in repositories if r.endswith(".repo") and os.path.isfile(r))

    cachedir = os.path.join(tempdir, "yum.cache")
    if not os.path.isdir(cachedir):
        os.mkdir(cachedir)

    logdir = os.path.join(tempdir, "yum.logs")
    if not os.path.isdir(logdir):
        os.mkdir(logdir)

    installroot = os.path.join(tempdir, "installroot")
    if not os.path.isdir(installroot):
        os.mkdir(installroot)

    yumconf = os.path.join(tempdir, "yum.conf")
    c = ConfigParser.ConfigParser()

    # add the main section
    section = "main"
    data = {"cachedir": cachedir,
            "keepcache": 0,
            "gpgcheck": 0,
            "plugins": 0,
            "assumeyes": 1,
            "reposdir": "",
            "tsflags": "nodocs"}

    if proxy:
        data["proxy"] = proxy

    c.add_section(section)
    map(lambda (key, value): c.set(section, key, value), data.items())

    # add the main repository - the first repository from list
    # This list may be empty if using .repo files
    if repo_urls:
        section = "lorax-repo"
        data = {"name": "lorax repo",
                "baseurl": repo_urls[0],
                "enabled": 1}

        c.add_section(section)
        map(lambda (key, value): c.set(section, key, value), data.items())

    # add the extra repositories
    for n, extra in enumerate(repo_urls[1:], start=1):
        section = "lorax-extra-repo-{0:d}".format(n)
        data = {"name": "lorax extra repo {0:d}".format(n),
                "baseurl": extra,
                "enabled": 1}

        c.add_section(section)
        map(lambda (key, value): c.set(section, key, value), data.items())

    # write the yum configuration file
    with open(yumconf, "w") as f:
        c.write(f)

    # create the yum base object
    yb = yum.YumBase()

    yb.preconf.fn = yumconf
    yb.preconf.root = installroot
    if releasever:
        yb.preconf.releasever = releasever

    # Turn on as much yum logging as we can
    yb.preconf.debuglevel = 10
    yb.preconf.errorlevel = 10
    yb.logger.setLevel(logging.DEBUG)
    yb.verbose_logger.setLevel(logging.DEBUG)

    # Add .repo files from the cmdline
    for fn in repo_files:
        if os.path.exists(fn):
            print("Adding repos from file: %s" % fn)
            yb.getReposFromConfigFile(fn)

    # Get the metadata
    print("Downloading metadata")
    for r in yb.repos.sort():
        r.metadata_expire = 0
        r.mdpolicy = "group:all"
    try:
        yb.doRepoSetup()
        yb.repos.doSetup()
        yb.repos.populateSack(mdtype='all', cacheonly=0)
    except YumBaseError as e:
        print("Failed to update metadata: %s" % str(e))
        sys.exit(1)

    return yb


def mako_installpkg(pkg_file):
    """ Read a Mako template file and return a list of all the packages
        listed in the installpkg lines.
    """
    packages = []
    with open(pkg_file, "r") as f:
        for line in f.readlines():
            if line.startswith("installpkg"):
                packages += line.split()[1:]

    return packages

def get_deps(ybo, packages):
    """Return the dependencies for the listed packages"""
    ybo.closeRpmDB()
    for pkg in packages:
        print("Adding %s to the transaction" % pkg)
        try:
            if pkg[0] == "@":
                ybo.selectGroup(pkg[1:], group_package_types=['mandatory', 'default'])
            else:
                ybo.install(pattern=pkg)
        except Exception as e:
            print("Failed to install %s\n%s" % (pkg, e))

    try:
        print("Checking dependencies")
        (rc, msg) = ybo.buildTransaction()
        if rc not in [0, 1, 2]:
            print("There was a problem depsolving %s: %s" % (packages, msg))
            return []
        ybo.tsInfo.makelists()
        return sorted([tm.po.nevra for tm in ybo.tsInfo.installed + ybo.tsInfo.depinstalled])
    except YumBaseError as e:
        print("Dependency check failed: %s" % e)
        return []

def find_package_parent(pkg_name, ybo, packages):
    """ Brute force discovery of what make a specific pkg_name get pulled
        into the transaction.
    """
    # Add packages to the set one at a time and examine the deps
    package_set = []
    for p in packages:
        package_set += [p]
        deps = get_deps(ybo, package_set)

        # Is the needle in this haystack?
        if any(pkg_name in pkg for pkg in deps):
            print("FOUND IT! %s pulled in %s" % (p, pkg_name))
            print(package_set)
            return 0
    return 1

if __name__ == "__main__":
    parser = setup_argparse()
    opts = parser.parse_args()

    tempdir = opts.tempdir or tempfile.mkdtemp(prefix="test-yum.")
    print("Using tempdir: %s" % tempdir)
    ybo = get_ybo(tempdir, opts.source, opts.release, opts.proxy)

    packages = opts.packages
    if opts.pkg_file:
        packages += mako_installpkg(opts.pkg_file)

    if len(packages) == 0:
        print("Pass packages on cmdline or via Mako template using -t")
        sys.exit(1)

    # Find what pulls in a specific package
    if opts.find:
        sys.exit(find_package_parent(opts.find, ybo, packages))

    deps = get_deps(ybo, packages)
    print("%d packages selected" % len(deps))
    if len(deps) == 0:
        print("No packages in transaction")

    # Print what Yum picked.
    for pkg in deps:
        print(pkg)