#!/usr/bin/python2

import argparse
import atexit
import os
import pwd
import subprocess
import shutil
import sys
import tempfile

sys.path.insert(0, "/usr/lib/entropy/entropy_path_loader")

import entropy_path_loader

from entropy.server.interfaces import Server

import entropy.dep


if __name__ == "__main__":

    srv = Server()

    def ValidString(arg):
        if not arg:
            raise ValueError("invalid string")
        return arg

    def ValidRepository(arg):
        if arg not in srv.repositories():
            raise ValueError("invalid repository %s" % (arg,))
        return arg

    parser = argparse.ArgumentParser(
        description="Bump Kernel Packages for Entropy Server")

    parser.add_argument(
        "query_version", metavar="<query version>",
        help="the uname version to query",
        type=ValidString)
    parser.add_argument(
        "build_version", metavar="<build version>",
        help="the uname version to build packages for",
        type=ValidString)

    parser.add_argument(
        "build_repo", metavar="<build repo>",
        help="the repository to place packages into",
        type=ValidRepository)

    parser.add_argument(
        "--only-injected", action="store_true", default=False,
        help="only consider injected packages")
    parser.add_argument(
        "--no-spm-repo", action="store_true", default=False,
        help=("do not include the SPM repository suffix "
              "in dependency atoms"))
    parser.add_argument(
        "--keyslot", action="store_true", default=False,
        help=("do not use the exact package version but "
              "rather key:slot"))
    parser.add_argument(
        "--canfail", action="store_true", default=False,
        help="if emerge can fail and you want to ignore it")

    parser.add_argument(
        "--non-interactive", action="store_true", default=False,
        help="if set, the script will run in non-interactive mode")

    parser.add_argument(
        "--disable-shell-wrap", action="store_true", default=False,
        help="if set, the script will not source .bashrc")

    nsargs = parser.parse_args(sys.argv[1:])

    known_multiple_injected_versions = set([
            "x11-drivers/nvidia-drivers:0",
            "x11-drivers/ati-drivers:0",
        ])

    pkgs_map = {}
    try:

        pkg_matches = {}
        repository_ids = srv.repositories()
        for repository_id in repository_ids:
            repo = srv.open_repository(repository_id)
            pkg_ids = repo.searchTaggedPackages(nsargs.query_version)

            for pkg_id in pkg_ids:
                key_slot = repo.retrieveKeySlotAggregated(pkg_id)
                obj = pkg_matches.setdefault(key_slot, set())
                obj.add((pkg_id, repository_id))

        injected_pkgs = []
        normal_pkgs = []
        for key_slot, candidates in pkg_matches.items():
            candidate = srv.atom_match(key_slot + "#" + nsargs.query_version)
            if candidate not in candidates:
                print("%s (%s) not in candidates %s, skipping" % (
                        candidate, key_slot, candidates))
                continue
            pkg_id, repository_id = candidate
            repo = srv.open_repository(repository_id)
            if repo.isInjected(pkg_id):
                injected_pkgs.append(candidate)
            else:
                normal_pkgs.append(candidate)

            # address other injected pkgs, like older nvidia drivers
            for other_candidate in candidates:
                if other_candidate != candidate:
                    pkg_id, repository_id = other_candidate
                    repo = srv.open_repository(repository_id)
                    if repo.isInjected(pkg_id):
                        clean_key_slot = entropy.dep.remove_tag_from_slot(
                            key_slot)
                        # filter out packages that we know we don't handle multiple
                        # versions of them.
                        if clean_key_slot in known_multiple_injected_versions:
                            injected_pkgs.append(other_candidate)

        normal_atoms = []
        injected_atoms = []
        data = (
            (normal_pkgs, normal_atoms),
            (injected_pkgs, injected_atoms)
        )

        spm = srv.Spm()

        for lst, dst in data:
            for pkg_id, repository_id in lst:
                repo = srv.open_repository(repository_id)
                if nsargs.keyslot:
                    atom_str = entropy.dep.remove_tag_from_slot(
                        repo.retrieveKeySlotAggregated(pkg_id))
                else:
                    atom_str = "~" + entropy.dep.remove_tag(
                        repo.retrieveAtom(pkg_id))
                if not nsargs.no_spm_repo:
                    spm_repo = repo.retrieveSpmRepository(pkg_id)
                    if spm_repo is not None:
                        atom_str += "::" + spm_repo

                if not spm.match_package(atom_str):
                    print("%s not found in Portage, skipping" % (atom_str,))
                    continue

                dst.append(atom_str)
    finally:
        srv.shutdown()

    if not nsargs.only_injected:
        if normal_atoms:
            print("normal packages: %s" % (
                " ".join(normal_atoms),))
    if injected_atoms:
        print("injected packages: %s" % (
            " ".join(injected_atoms),))

    if not (normal_atoms or injected_atoms):
        print("nothing to do !!")
        raise SystemExit(0)

    def clean_dirs(xdirs):
        for directory in xdirs:
            try:
                shutil.rmtree(directory, True)
            except shutil.Error:
                pass

    dirs = []
    atexit.register(clean_dirs, dirs)

    os.environ['KERNEL_DIR'] = "/usr/src/linux-%s" % (
        nsargs.build_version,)

    portage_tmpdir = tempfile.mkdtemp(
        dir="/var/tmp", prefix="bump_kernel_packages.portage_tmpdir")
    os.chmod(portage_tmpdir, 0o775)

    try:
        pass_s = pwd.getpwnam("portage")
        os.chown(portage_tmpdir, pass_s.pw_uid, pass_s.pw_gid)
    except KeyError:
        pass

    dirs.append(portage_tmpdir)
    os.environ["PORTAGE_TMPDIR"] = portage_tmpdir
    cmd_prefix = ""
    if not nsargs.disable_shell_wrap:
        cmd_prefix = "test -e ~/.bashrc && source ~/.bashrc;"

    if normal_atoms and not nsargs.only_injected:
        args = ["emerge", "-v", "--keep-going"]
        if not nsargs.non_interactive:
            args.append("-a")
        args += [x for x in normal_atoms]
        # emerge is a shell function in bashrc
        rc = subprocess.call(cmd_prefix + " ".join(args), shell=True)
        if rc != 0 and not nsargs.canfail:
            raise SystemExit(rc)
        subprocess.call(["/usr/sbin/etc-update"])

        args = ["eit", "add", "--to", nsargs.build_repo]
        if nsargs.non_interactive:
            args.append("--quick")

        args += normal_atoms
        # eit is a shell function in bashrc
        rc = subprocess.call(cmd_prefix + " ".join(args), shell=True)
        if rc != 0:
            raise SystemExit(rc)

    if injected_atoms:
        tmp_dir = tempfile.mkdtemp(
            dir="/var/tmp", prefix="bump_kernel_packages")
        dirs.append(tmp_dir)
        os.environ['PKGDIR'] = tmp_dir

        args = ["emerge", "-Bv", "--nodeps", "--keep-going"]
        if not nsargs.non_interactive:
            args.append("-a")
        args += [x for x in injected_atoms]
        # emerge is a shell function in .bashrc
        rc = subprocess.call(cmd_prefix + " ".join(args), shell=True)
        if rc != 0 and not nsargs.canfail:
            raise SystemExit(rc)

        tbz2s = []
        for category in os.listdir(tmp_dir):
            path = os.path.join(tmp_dir, category)
            if not os.path.isdir(path):
                continue
            for sub_file in os.listdir(path):
                if not sub_file.endswith(".tbz2"):
                    continue
                tbz2s.append(os.path.join(path, sub_file))

        if tbz2s:
            args = ["eit", "inject", "--to", nsargs.build_repo]
            args += tbz2s
            # eit is a shell function in bashrc
            rc = subprocess.call(cmd_prefix + " ".join(args), shell=True)
            if rc != 0:
                raise SystemExit(rc)

    raise SystemExit(0)
