#!/cvmfs/dirac.egi.eu/dirac/v8.0.48/Linux-x86_64/bin/python
# PYTHON_ARGCOMPLETE_OK

# Copyright 2012-2023, Andrey Kislyuk and argcomplete contributors.
# Licensed under the Apache License. See https://github.com/kislyuk/argcomplete for more info.

"""
Activate the generic bash-completion script or zsh completion autoload function for the argcomplete module.
"""

import argparse
import os
import shutil
import site
import subprocess
import sys

import argcomplete

zsh_shellcode = """
# Begin added by argcomplete
fpath=( {zsh_fpath} "${{fpath[@]}}" )
# End added by argcomplete
"""

bash_shellcode = """
# Begin added by argcomplete
source "{activator}"
# End added by argcomplete
"""


def get_local_dir():
    try:
        return subprocess.check_output(["brew", "--prefix"]).decode().strip()
    except (FileNotFoundError, subprocess.CalledProcessError):
        return "/usr/local"


def get_zsh_system_dir():
    return f"{get_local_dir()}/share/zsh/site-functions"


def get_bash_system_dir():
    if "BASH_COMPLETION_COMPAT_DIR" in os.environ:
        return os.environ["BASH_COMPLETION_COMPAT_DIR"]
    elif sys.platform == "darwin":
        return f"{get_local_dir()}/etc/bash_completion.d"  # created by homebrew
    else:
        return "/etc/bash_completion.d"  # created by bash-completion


def get_activator_dir():
    return os.path.join(os.path.abspath(os.path.dirname(argcomplete.__file__)), "bash_completion.d")


def get_activator_path():
    return os.path.join(get_activator_dir(), "_python-argcomplete")


def install_to_destination(dest):
    activator = get_activator_path()
    if dest == "-":
        with open(activator) as fh:
            sys.stdout.write(fh.read())
        return
    destdir = os.path.dirname(dest)
    if not os.path.exists(destdir):
        try:
            os.makedirs(destdir, exist_ok=True)
        except Exception as e:
            parser.error(f"path {destdir} does not exist and could not be created: {e}")
    try:
        print(f"Installing {activator} to {dest}...", file=sys.stderr)
        shutil.copy(activator, dest)
        print("Installed.", file=sys.stderr)
    except Exception as e:
        parser.error(
            f"while installing to {dest}: {e}. Please run this command using sudo, or see --help for more options."
        )


def get_consent():
    if args.yes is True:
        return True
    while True:
        res = input("OK to proceed? [y/n] ")
        if res.lower() not in {"y", "n", "yes", "no"}:
            print('Please answer "yes" or "no".', file=sys.stderr)
        elif res.lower() in {"y", "yes"}:
            return True
        else:
            return False


def append_to_config_file(path, shellcode):
    if os.path.exists(path):
        with open(path, 'r') as fh:
            if shellcode in fh.read():
                print(f"The code already exists in the file {path}.", file=sys.stderr)
                return
        print(f"argcomplete needs to append to the file {path}. The following code will be appended:", file=sys.stderr)
        for line in shellcode.splitlines():
            print(">", line, file=sys.stderr)
        if not get_consent():
            print("Not added.", file=sys.stderr)
            return
    print(f"Adding shellcode to {path}...", file=sys.stderr)
    with open(path, "a") as fh:
        fh.write(shellcode)
    print("Added.", file=sys.stderr)


def link_user_rcfiles():
    # TODO: warn if running as superuser
    zsh_rcfile = os.path.join(os.path.expanduser(os.environ.get("ZDOTDIR", "~")), ".zshenv")
    append_to_config_file(zsh_rcfile, zsh_shellcode.format(zsh_fpath=get_activator_dir()))

    bash_completion_user_file = os.path.expanduser("~/.bash_completion")
    append_to_config_file(bash_completion_user_file, bash_shellcode.format(activator=get_activator_path()))


parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("-y", "--yes", help="automatically answer yes for all questions", action="store_true")
parser.add_argument("--dest", help='Specify the shell completion modules directory to install into, or "-" for stdout')
parser.add_argument("--user", help="Install into user directory", action="store_true")
argcomplete.autocomplete(parser)
args = parser.parse_args()
destinations = []

if args.dest:
    if args.dest != "-" and not os.path.exists(args.dest):
        parser.error(f"directory {args.dest} was specified via --dest, but it does not exist")
    destinations.append(args.dest)
elif site.ENABLE_USER_SITE and site.USER_SITE in argcomplete.__file__:
    print(
        "Argcomplete was installed in the user site local directory. Defaulting to user installation.", file=sys.stderr
    )
    link_user_rcfiles()
elif sys.prefix != sys.base_prefix:
    print("Argcomplete was installed in a virtual environment. Defaulting to user installation.", file=sys.stderr)
    link_user_rcfiles()
elif args.user:
    link_user_rcfiles()
else:
    print("Defaulting to system-wide installation.", file=sys.stderr)
    destinations.append(f"{get_zsh_system_dir()}/_python-argcomplete")
    destinations.append(f"{get_bash_system_dir()}/python-argcomplete")

for destination in destinations:
    install_to_destination(destination)

if args.dest is None:
    print("Please restart your shell or source the installed file to activate it.", file=sys.stderr)