#!/usr/bin/env python3

import argparse
import sys
import os
import glob
try:
    from shlex import quote as shell_quote
except ImportError:
    from pipes import quote as shell_quote

# Set up PYTHONPATH to use the output of 'setup.py build', which will
# be in a subdirectory called 'lib.<something>' of the 'python-build'
# directory. The <something> will depend on the particular host
# architecture, so we use glob to find it, and hope to find only one
# of them.

def find_libdir(builddir):
    libdir_globs = [
        "lib.*-{major:d}.{minor:d}{dbgsuffix}".format(
            major = sys.version_info.major,
            minor = sys.version_info.minor,
            dbgsuffix = ("-pydebug" if (hasattr(sys, 'abiflags') and
                                        'd' in sys.abiflags) else "")),
        "lib.*-cpython-{major:d}{minor:d}{dbgsuffix}".format(
            major = sys.version_info.major,
            minor = sys.version_info.minor,
            dbgsuffix = ("-pydebug" if (hasattr(sys, 'abiflags') and
                                        'd' in sys.abiflags) else "")),
    ]

    libdirs = []
    for libdir_glob in libdir_globs:
        libdirs.extend(glob.glob(os.path.join(builddir, libdir_glob)))
    if len(libdirs) > 1:
        sys.stderr.write("Found multiple possible lib directories:" + "".join(
                ["\n  %s" % libdir for libdir in libdirs]) + "\n")
        sys.exit(1)
    if len(libdirs) == 0:
        sys.stderr.write("Found no possible lib directory in %s\n" % builddir)
        sys.exit(1)

    return libdirs[0]

def make_pythonpath(libdir):
    oldpath = os.environ.get("PYTHONPATH", "")
    if len(oldpath) > 0:
        return libdir + ":" + oldpath
    else:
        return libdir

def run_python(pythonpath, args):
    os.environ["PYTHONPATH"] = pythonpath
    args = [os.path.basename(sys.executable)] + args
    os.execvp(sys.executable, args)

def run_program(pythonpath, args):
    os.environ["PYTHONPATH"] = pythonpath
    os.execvp(args[0], args)

def print_environment(pythonpath):
    sys.stdout.write("PYTHONPATH={}; export PYTHONPATH\n".format(
        shell_quote(pythonpath)))

def main():
    parser = argparse.ArgumentParser(
        description='Set up PYTHONPATH to run out of a Python distutils '
        'build directory.')
    parser.add_argument("--build-base", help="Build base directory as given"
                        " to setup.py.")
    parser.add_argument("command", nargs=argparse.REMAINDER,
                        help='Arguments to run Python interpreter with.')
    parser.add_argument("--exec", action='store_true',
                        help="Execute a command other than Python.")
    args = parser.parse_args()

    builddir = args.build_base or os.path.join(
        os.path.dirname(os.path.abspath(__file__)), "..", "python-build")
    builddir = os.path.abspath(builddir)
    libdir = find_libdir(builddir)
    pythonpath = make_pythonpath(libdir)

    if len(args.command) > 0:
        if args.command[0] == "--":
            args.command = args.command[1:]
        if getattr(args, 'exec'):
            if len(args.command) < 1:
                sys.exit("testenv: must provide a command name with --exec")
            run_program(pythonpath, args.command)
        else:
            run_python(pythonpath, args.command)
    else:
        print_environment(pythonpath)

if __name__ == '__main__':
    main()
