#!/usr/bin/env python
"""Script to install bundled Ansible."""

from __future__ import print_function

import argparse
import os
import shutil
import subprocess
import sys
import tarfile
import tempfile
from contextlib import contextmanager

DEFAULT_INSTALL_DIR = "/opt/ansible"
PACKAGES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "packages")
ANSIBLE_BINS = [
    "ansible",
    "ansible-config",
    "ansible-connection",
    "ansible-console",
    "ansible-doc",
    "ansible-galaxy",
    "ansible-inventory",
    "ansible-playbook",
    "ansible-pull",
    "ansible-test",
    "ansible-vault",
]
GTE_PY35 = sys.version_info[:2] >= (3, 5)
UNSUPPORTED_PYTHON = (sys.version_info[0] == 2 and sys.version_info[:2] <= (2, 6)) or (
    sys.version_info[0] == 3 and sys.version_info[:2] <= (3, 4)
)


def fail(msg, exit_code=1):
    """Exit script with error and message."""
    sys.stderr.write("{0}\n".format(msg))
    sys.exit(exit_code)


@contextmanager
def cd(dirname):
    """Change into dirname directory."""
    original = os.getcwd()
    os.chdir(dirname)
    try:
        yield
    finally:
        os.chdir(original)


def run(cmd):
    """Execute a sub-process."""
    print("Running cmd: {}".format(cmd))
    p = subprocess.Popen(
        cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    stdout, stderr = p.communicate()
    if p.returncode != 0:
        output = (stdout + stderr).decode("utf-8")
        fail("Bad rc (%s) for cmd '%s': %s" % (p.returncode, cmd, output))
    return stdout


def create_install_structure(working_dir, install_dir):
    """Create installation directory structure and virtualenv."""
    if not os.path.isdir(install_dir):
        os.makedirs(install_dir)
    _create_virtualenv(location=install_dir, working_dir=working_dir)


def _create_virtualenv(location, working_dir):
    """Create virtualenv for installation."""
    if GTE_PY35:
        # use internal venv module
        run("{} -m venv {}".format(sys.executable, location))
    else:
        # use external virtualenv package
        venv_package = os.path.join(
            PACKAGES_DIR, _get_package_tarball(PACKAGES_DIR, "virtualenv")
        )

        with tarfile.open(venv_package) as tar:
            tar.extractall(path=working_dir)
            # virtualenv should be the only directory, so we can pick the first one
            with cd(os.path.join(working_dir, os.listdir(working_dir)[0])):
                run(
                    "{0} virtualenv.py --no-download {1}".format(
                        sys.executable, location
                    )
                )


def _get_package_tarball(package_dir, package_prefix):
    package_filenames = sorted(
        [p for p in os.listdir(package_dir) if p.startswith(package_prefix)]
    )
    return package_filenames[-1]


def pip_install_packages(install_dir):
    """Install tarball pip packages."""
    ansible_tarball = [p for p in os.listdir(PACKAGES_DIR) if p.startswith("ansible")]
    if len(ansible_tarball) != 1:
        message = (
            "Multiple versions of Ansible were found in %s. Please clear "
            "out this directory before proceeding."
        )
        fail(message % PACKAGES_DIR)
    # install wheel first to allow other installs to be successful
    _pip_install_package(install_dir, "wheel")
    _pip_install_package(install_dir, "ansible")


def _pip_install_package(install_dir, package):
    pip_script = os.path.join(install_dir, "bin", "pip")
    run(
        "%s install --no-cache-dir --no-index --find-links %s %s"
        % (pip_script, PACKAGES_DIR, package)
    )


def create_symlinks(real_dir, link_dir):
    """Create symlinks to each Ansible binary."""
    result = True
    for cmd in ANSIBLE_BINS:
        link_path = os.path.join(link_dir, cmd)
        if os.path.isfile(link_path) and not os.path.islink(link_path):
            print("File {} already exists and is not a symlink".format(link_path))
            result = False
    if result:
        if not os.path.isdir(link_dir):
            os.makedirs(link_dir)
        for cmd in ANSIBLE_BINS:
            real_path = os.path.join(real_dir, cmd)
            link_path = os.path.join(link_dir, cmd)
            if os.path.isfile(link_path):
                print("Symlink already exists: {}".format(link_path))
                print("Removing symlink.")
                os.remove(link_path)
            os.symlink(real_path, link_path)
    return result


if UNSUPPORTED_PYTHON:
    fail(
        "Unsupported Python version %s.%s\n"
        "Must be Python 2.7+ or Python 3.5+" % sys.version_info[:2]
    )


def main():
    """Execute."""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-i",
        "--install-dir",
        help="The directory to install Ansible. The default is {}".format(
            DEFAULT_INSTALL_DIR
        ),
        default=DEFAULT_INSTALL_DIR,
    )
    parser.add_argument(
        "-l",
        "--link-dir",
        metavar="DIR",
        help="If this argument is provided, symlinks will be created in this "
        "directory for all the Ansible executables. "
        "Useful if you want Ansible executables somewhere in your path, "
        "e.g. -l /usr/local/bin.",
    )
    args = parser.parse_args()

    if args.link_dir and not os.path.isdir(args.link_dir):
        fail("--link-dir {} provided but does not exist".format(args.link_dir))
    working_dir = tempfile.mkdtemp()
    try:
        install_dir = os.path.abspath(args.install_dir)
        create_install_structure(working_dir, install_dir)
        pip_install_packages(install_dir)
        real_dir = os.path.join(install_dir, "bin")
        if args.link_dir:
            result = create_symlinks(real_dir, args.link_dir)
            if not result:
                print("Symlinks not created")
        print("Installation complete.")
    finally:
        shutil.rmtree(working_dir)


if __name__ == "__main__":
    main()
