#!/usr/bin/env python3
"""List pip dependencies or system package dependencies for cloud-init."""

# You might be tempted to rewrite this as a shell script, but you
# would be surprised to discover that things like 'egrep' or 'sed' may
# differ between Linux and *BSD.

try:
    from argparse import ArgumentParser
except ImportError:
    raise RuntimeError(
        "Could not import argparse. Please install python3-argparse "
        "package to continue"
    )

import json
import os
import re
import subprocess
import sys

DEFAULT_REQUIREMENTS = "requirements.txt"

# Map the appropriate package dir needed for each distro choice
DISTRO_PKG_TYPE_MAP = {
    "centos": "redhat",
    "eurolinux": "redhat",
    "miraclelinux": "redhat",
    "fedora": "fedora",
    "rocky": "redhat",
    "redhat": "redhat",
    "debian": "debian",
    "ubuntu": "debian",
    "opensuse": "suse",
    "opensuse-leap": "suse",
    "opensuse-microos": "suse",
    "opensuse-tumbleweed": "suse",
    "sle_hpc": "suse",
    "sle-micro": "suse",
    "sles": "suse",
    "suse": "suse",
}

MAYBE_RELIABLE_YUM_INSTALL = [
    "sh",
    "-c",
    """
    error() { echo "$@" 1>&2; }
    configure_repos_for_proxy_use() {
        if [ -f /etc/yum.conf ]; then
            CFG_FILE=/etc/yum.conf
        fi
        if [ -f /etc/dnf/dnf.conf ]; then
            CFG_FILE=/etc/dnf/dnf.conf
        fi
        if [ ! "${CFG_FILE}" ]; then
            error "No yum.conf or dnf.conf"
            return 0
        fi
        grep -q "^proxy=" $CFG_FILE || return 0
        error ":: http proxy in use => forcing the use of fixed URLs in /etc/yum.repos.d/*.repo"
        sed -i --regexp-extended '/^#baseurl=/s/#// ; /^(mirrorlist|metalink)=/s/^/#/' /etc/yum.repos.d/*.repo
        sed -i 's/download\\.fedoraproject\\.org/dl.fedoraproject.org/g' /etc/yum.repos.d/*.repo
        sed -i 's/download\\.example/dl.fedoraproject.org/g' /etc/yum.repos.d/*.repo
        CENTOS_REPO="/etc/yum.repos.d/centos.repo"
        if [ -f $CENTOS_REPO ]; then
            grep -q baseurl $CENTOS_REPO
            if [ $? -eq 1 ]; then
                # CentOS 9 does not provide baseurl definitions
                sed -i '/\\[baseos\\]/a baseurl=https://mirror.stream.centos.org/9-stream/BaseOS/x86_64/os' ${CENTOS_REPO}
                sed -i '/\\[appstream\\]/a baseurl=https://mirror.stream.centos.org/9-stream/AppStream/x86_64/os' ${CENTOS_REPO}
                sed -i '/\\[crb\\]/a baseurl=https://mirror.stream.centos.org/9-stream/CRB/x86_64/os' ${CENTOS_REPO}
            fi
            CENTOS_EXTRAS_REPO="/etc/yum.repos.d/centos-addons.repo"
            if [ -f $CENTOS_EXTRAS_REPO ]; then
                sed -i '/\\[extras-common\\]/a baseurl=https://mirror.stream.centos.org/SIGs/9-stream/extras/x86_64/extras-common' ${CENTOS_EXTRAS_REPO}
            fi
        fi
    }
    configure_repos_for_proxy_use
    n=0; max=10;
    bcmd="yum install --downloadonly --assumeyes --setopt=keepcache=1"
    while n=$(($n+1)); do
       error ":: running $bcmd $* [$n/$max]"
       $bcmd "$@"
       r=$?
       [ $r -eq 0 ] && break
       [ $n -ge $max ] && { error "gave up on $bcmd"; exit $r; }
       nap=$(($n*5))
       error ":: failed [$r] ($n/$max). sleeping $nap."
       sleep $nap
    done
    error ":: running yum install --cacheonly --assumeyes $*"
    yum install --cacheonly --assumeyes "$@"
    configure_repos_for_proxy_use
    """,  # noqa: E501
    "reliable-yum-install",
]

ZYPPER_INSTALL = [
    "zypper",
    "--non-interactive",
    "--gpg-auto-import-keys",
    "install",
    "--auto-agree-with-licenses",
]

DRYRUN_DISTRO_INSTALL_PKG_CMD = {
    "redhat": ["yum", "install", "--assumeyes"],
    "fedora": ["yum", "install", "--assumeyes"],
}

DISTRO_INSTALL_PKG_CMD = {
    "redhat": MAYBE_RELIABLE_YUM_INSTALL,
    "fedora": MAYBE_RELIABLE_YUM_INSTALL,
    "debian": ["apt", "install", "-y"],
    "suse": ZYPPER_INSTALL,
}

DISTRO_UPDATE_PKG_CMD = {
    "redhat": ["yum", "update"],
    "debian": ["apt", "update", "-q"],
    "suse": ["zypper", "update"],
}

# List of base system packages required to enable ci automation
CI_SYSTEM_BASE_PKGS = {
    "common": ["make", "sudo", "tar"],
    "eurolinux": ["python3-tox"],
    "miraclelinux": ["python3-tox"],
    "fedora": ["python3-tox"],
    "redhat": ["python3-tox"],
    "centos": ["python3-tox"],
    "ubuntu": ["devscripts", "python3-dev", "libssl-dev", "tox", "sbuild"],
    "debian": ["devscripts", "python3-dev", "libssl-dev", "tox", "sbuild"],
}


# JSON definition of distro-specific package dependencies
DISTRO_PKG_DEPS_PATH = "packages/pkg-deps.json"


def get_parser():
    """Return an argument parser for this command."""
    parser = ArgumentParser(description=__doc__)
    parser.add_argument(
        "-r",
        "--requirements-file",
        type=str,
        dest="req_files",
        action="append",
        default=None,
        help="pip-style requirements file [default=%s]" % DEFAULT_REQUIREMENTS,
    )
    parser.add_argument(
        "-d",
        "--distro",
        type=str,
        choices=DISTRO_PKG_TYPE_MAP.keys(),
        help="The name of the distro to generate package deps for.",
    )
    deptype = parser.add_mutually_exclusive_group()
    deptype.add_argument(
        "-R",
        "--runtime-requires",
        action="store_true",
        default=False,
        dest="runtime_requires",
        help="Print only runtime required packages",
    )
    deptype.add_argument(
        "-b",
        "--build-requires",
        action="store_true",
        default=False,
        dest="build_requires",
        help="Print only buildtime required packages",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        default=False,
        dest="dry_run",
        help="Dry run the install, making no package changes.",
    )
    parser.add_argument(
        "-s",
        "--system-pkg-names",
        action="store_true",
        default=False,
        dest="system_pkg_names",
        help="Generate distribution package names (python3-pkgname).",
    )
    parser.add_argument(
        "-i",
        "--install",
        action="store_true",
        default=False,
        dest="install",
        help="When specified, install the required system packages.",
    )
    parser.add_argument(
        "-t",
        "--test-distro",
        action="store_true",
        default=False,
        dest="test_distro",
        help="Additionally install continuous integration system packages "
        "required for build and test automation.",
    )
    return parser


def get_package_deps_from_json(topdir, distro):
    """Get a dict of build and runtime package requirements for a distro.

    @param topdir: The root directory in which to search for the
        DISTRO_PKG_DEPS_PATH json blob of package requirements information.
    @param distro: The specific distribution shortname to pull dependencies
        for.
    @return: Dict containing "requires", "build-requires" and "rename" lists
        for a given distribution.
    """
    with open(os.path.join(topdir, DISTRO_PKG_DEPS_PATH), "r") as stream:
        deps = json.loads(stream.read())
    if distro is None:
        return {}
    if deps.get(distro):  # If we have a specific distro defined, use it.
        return deps[distro]
    # Use generic distro dependency map via DISTRO_PKG_TYPE_MAP
    return deps[DISTRO_PKG_TYPE_MAP[distro]]


def parse_pip_requirements(requirements_path):
    """Return the pip requirement names from pip-style requirements_path."""
    dep_names = []
    with open(requirements_path, "r") as fp:
        for line in fp:
            line = line.strip()
            if not line or line.startswith(("#", "-r ")):
                continue

            # remove pip-style markers
            dep = line.split(";")[0]

            # remove version requirements
            version_comparison = re.compile(r"[~!>=.<]+")
            if version_comparison.search(dep):
                dep_names.append(version_comparison.split(dep)[0].strip())
            else:
                dep_names.append(dep)
    return dep_names


def translate_pip_to_system_pkg(pip_requires, renames):
    """Translate pip package names to distro-specific package names.

    @param pip_requires: List of versionless pip package names to translate.
    @param renames: Dict containing special case renames from pip name to
        system package name for the distro.
    """
    prefix = "python3-"
    standard_pkg_name = "{0}{1}"
    translated_names = []
    for pip_name in pip_requires:
        pip_name = pip_name.lower()
        # Find a rename if present for the distro package and python version
        rename = renames.get(pip_name, "")
        if rename:
            translated_names.append(rename)
        else:
            translated_names.append(standard_pkg_name.format(prefix, pip_name))
    return translated_names


def main(distro):
    parser = get_parser()
    args = parser.parse_args()
    if "CLOUD_INIT_TOP_D" in os.environ:
        topd = os.path.realpath(os.environ.get("CLOUD_INIT_TOP_D"))
    else:
        topd = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))

    if args.test_distro:
        # Give us all the system deps we need for continuous integration
        if args.req_files:
            sys.stderr.write(
                "Parameter --test-distro overrides --requirements-file. Use "
                "one or the other.\n"
            )
            sys.exit(1)
        args.req_files = [
            os.path.join(topd, DEFAULT_REQUIREMENTS),
            os.path.join(topd, "test-" + DEFAULT_REQUIREMENTS),
        ]
        args.install = True
    if args.req_files is None:
        args.req_files = [os.path.join(topd, DEFAULT_REQUIREMENTS)]
        if not os.path.isfile(args.req_files[0]):
            sys.stderr.write(
                "Unable to locate '%s' file that should "
                "exist in cloud-init root directory." % args.req_files[0]
            )
            sys.exit(1)

    bad_files = [r for r in args.req_files if not os.path.isfile(r)]
    if bad_files:
        sys.stderr.write(
            "Unable to find requirements files: %s\n" % ",".join(bad_files)
        )
        sys.exit(1)

    pip_pkg_names = set()
    for req_path in args.req_files:
        pip_pkg_names.update(set(parse_pip_requirements(req_path)))
    deps_from_json = get_package_deps_from_json(topd, args.distro)
    renames = deps_from_json.get("renames", {})
    translated_pip_names = translate_pip_to_system_pkg(pip_pkg_names, renames)
    all_deps = []
    select_requires = [args.build_requires, args.runtime_requires]
    if args.distro:
        if not any(select_requires):
            all_deps.extend(
                translated_pip_names
                + deps_from_json["requires"]
                + deps_from_json["build-requires"]
            )
        else:
            if args.build_requires:
                all_deps.extend(deps_from_json["build-requires"])
            else:
                all_deps.extend(
                    translated_pip_names + deps_from_json["requires"]
                )
    else:
        if args.system_pkg_names:
            all_deps = translated_pip_names
        else:
            all_deps = pip_pkg_names
    all_deps = sorted(all_deps)
    if args.install:
        pkg_install(all_deps, args.distro, args.test_distro, args.dry_run)
    else:
        print("\n".join(all_deps))


def pkg_install(pkg_list, distro, test_distro=False, dry_run=False):
    """Install a list of packages using the DISTRO_INSTALL_PKG_CMD."""
    if test_distro:
        pkg_list = list(pkg_list) + CI_SYSTEM_BASE_PKGS["common"]
        distro_base_pkgs = CI_SYSTEM_BASE_PKGS.get(distro, [])
        pkg_list += distro_base_pkgs
    print(
        "Installing deps: {0}{1}".format(
            "(dryrun)" if dry_run else "", " ".join(pkg_list)
        )
    )
    distro_family = DISTRO_PKG_TYPE_MAP[distro]
    install_cmd = []
    update_cmd = DISTRO_UPDATE_PKG_CMD.get(distro_family, [])
    if dry_run:
        install_cmd.append("echo")
        if update_cmd:
            update_cmd.insert(0, "echo")
    if os.geteuid() != 0:
        install_cmd.append("sudo")
        if update_cmd:
            update_cmd.append("sudo")

    if dry_run and distro_family in DRYRUN_DISTRO_INSTALL_PKG_CMD:
        cmd = DRYRUN_DISTRO_INSTALL_PKG_CMD[distro_family]
    else:
        cmd = DISTRO_INSTALL_PKG_CMD[distro_family]
    install_cmd.extend(cmd)

    if update_cmd:
        subprocess.check_call(update_cmd)

    if distro in ["centos", "redhat", "rocky", "eurolinux"]:
        # CentOS and Redhat need epel-release to access oauthlib and jsonschema
        subprocess.check_call(install_cmd + ["epel-release"])
        subprocess.check_call(
            ["dnf", "config-manager", "--set-enabled", "crb"]
        )
        # Disable epel-cisco-openh264 if present, we don't need ffmeg codex
        # and additional repos just puts unnecessary load on our proxies
        subprocess.call(
            ["dnf", "config-manager", "--set-disabled", "epel-cisco-openh264"]
        )
    if DISTRO_PKG_TYPE_MAP.get(distro) in ["redhat", "suse", "fedora"]:
        pkg_list.append("rpm-build")
    subprocess.check_call(install_cmd + pkg_list)


if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()
    sys.exit(main(args.distro))