# Copyright (c) 2009-2015 by Enthought, Inc.
# All rights reserved.

# This script is run after core Python was installed on Unix systems.
# It does the following:
#   * update the first line in scripts in sys.prefix/bin (sys.prefix/Scripts on
#   windows) with sys.executable
#   * update the prefix=... line in config/Makefile with sys.prefix
#   * generate activate scripts

import errno
import io
import os
import os.path
import struct
import sys
import re

import venv


hashbang_pat = re.compile(r'#!(.+)$', re.M)

verbose = False


MAX_PATH = 260
# Max size of a comment appended to a COFF exe
MAX_COMMENT_SIZE = 2 ** 16


class LauncherSignatureNotFound(Exception):
    pass


# Signature used in new 'simple launcher' as used in entry points created by
# from wheel (e.g. in versions of pip >= 7.0)
SIG = b"\x50\x4b\x05\x06"


def find_shebang_pos(path):
    """ Look for shebang start/end position for the given path.

    This works for new style launchers (as created by pypa/distlib 0.2.0 and
    after).

    Raises a LauncherSignatureNotFound exception if a valid signature cannot be
    found.
    """
    with open(path, 'rb') as fp:
        fp.seek(0, 2)
        size = fp.tell()

        n = MAX_COMMENT_SIZE
        pos = size - n
        if pos < 0:
            pos = 0
        fp.seek(pos, 0)

        data = fp.read(n)

    relative_sig_pos = data.find(SIG)
    if relative_sig_pos == -1:
        raise LauncherSignatureNotFound(
            "No signature found in {!r}".format(path)
        )

    end_cdr = struct.unpack("LLLLL", data[relative_sig_pos:relative_sig_pos+20])
    cdsize = end_cdr[3]
    cdoffset = end_cdr[4]

    end_cdr_offset = relative_sig_pos - (cdsize + cdoffset)

    shebang_start = end_cdr_offset - MAX_PATH
    assert shebang_start > 0

    shebang_start = data.find(b"#!", shebang_start, end_cdr_offset)
    shebang_end = end_cdr_offset

    return pos + shebang_start, pos + shebang_end


def copy_until(source, target, until, bufsize):
    """Copy until bytes from source into target."""
    nblocks = until % bufsize
    remain = until // bufsize

    for i in range(nblocks):
        target.write(source.read(bufsize))
    target.write(source.read(remain))


def copy_until_end(source, target, bufsize):
    """Copy content from source into target until source EOF."""
    while True:
        buf = source.read(bufsize)
        if buf:
            target.write(buf)
        else:
            break


def rename(source, target):
    if os.name == "nt":
        try:
            os.rename(source, target)
        except WindowsError as e:
            if e.errno != errno.EEXIST:
                raise
            else:
                os.unlink(target)
                os.rename(source, target)
    else:
        return os.rename(source, target)


def fix_new_launcher(path):
    shebang_start, shebang_end = find_shebang_pos(path)

    new_shebang_line = (
        "#!" + sys.executable.replace("\\", "\\\\") + "\n\r\n"
    ).encode("utf8")

    tmpfile = path + ".tmp"

    bufsize = 65384

    with open(path, "rb") as fpin:
        try:
            with open(tmpfile, "wb") as fpout:
                copy_until(fpin, fpout, shebang_start, bufsize)
                fpout.write(new_shebang_line)
                fpin.seek(shebang_end, 0)
                copy_until_end(fpin, fpout, bufsize)
        except IOError:
            os.unlink(tmpfile)
            raise

    rename(tmpfile, path)


def makedirs(p):
    try:
        os.makedirs(p)
    except OSError as e:
        if e.errno != errno.EEXIST:
            raise


def read_data(path):
    with open(path) as fi:
        return fi.read()


def write_data(path, data):
    if verbose:
        print "Updating: %r" % path
    with open(path, 'w') as fo:
        fo.write(data)


def fix_script(path):
    if os.path.islink(path) or not os.path.isfile(path):
        return

    data = read_data(path)
    m = hashbang_pat.match(data)
    if not (m and 'python' in m.group(1).lower()):
        return

    replacement = "#!{0}".format(sys.executable.replace("\\", "\\\\"))
    new_data = hashbang_pat.sub(replacement, data, count=1)
    if new_data == data:
        return
    write_data(path, new_data)
    os.chmod(path, 0755)


def fix_makefile(path):
    if verbose:
        print 'Fixing Mafefile: %r' % path
    data = data_org = read_data(path)

    # set the current prefix
    pre_pat = re.compile(r'^prefix=.*$', re.M)
    data = pre_pat.sub('prefix=\t\t' + sys.prefix, data)

    # replace the placeholders with nothing
    plc_pat = re.compile(r'-R(/PLACEHOLD){5,}')
    data = plc_pat.sub('', data)

    # make sure the build include directories are not listed in OPT=
    opt_pat = re.compile(r'^OPT=.*$', re.M)
    m = opt_pat.search(data)
    if m:
        line = m.group()
        line = re.sub(r'-I\S+', '', line).rstrip()
        data = opt_pat.sub(line, data)
    else:
        print "WARNING: Did not find 'OPT=...' in %r" % path

    # fix LDFLAGS=
    ldf_pat = re.compile(r'^LDFLAGS=.*$', re.M)
    m = ldf_pat.search(data)
    if m:
        line = m.group()
        line = re.sub(r'-L\S+', '', line).rstrip()
        data = ldf_pat.sub(line, data)
    else:
        print "WARNING: Did not find 'LDFLAGS=...' in %r" % path

    # see if the data was actually changed
    if data == data_org:
        if verbose:
            print "Hmm: Already up-to-date: %r" % path
        return
    write_data(path, data)


def fix_remove(path):
    data = data_org = read_data(path)
    data = data.replace('PREFIX=`dirname $_SCRIPT_DIR`',
                        'PREFIX=%s' % sys.prefix)
    if data == data_org:
        if verbose:
            print "Hmm: Already up-to-date: %r" % path
        return
    write_data(path, data)
    os.chmod(path, 0755)


def create_activation_scripts():
    venv_pkg_dir = os.path.normpath(os.path.dirname(venv.__file__))
    venv_scripts_dir = os.path.join(venv_pkg_dir, "scripts", os.name)

    clean_exec_prefix = os.path.normpath(os.path.abspath(sys.exec_prefix))
    prompt_name = os.path.basename(clean_exec_prefix)

    prompt = "({0})".format(prompt_name)
    venv_dir = sys.prefix
    if sys.platform == 'win32':
        venv_bin_name = "Scripts"
    else:
        venv_bin_name = "bin"
    scriptsdir = os.path.join(venv_dir, venv_bin_name)

    makedirs(scriptsdir)

    for root, dirs, files in os.walk(venv_scripts_dir):
        for f in files:
            source = os.path.join(root, f)
            target = os.path.join(scriptsdir, os.path.basename(source))

            with open(source, "rt") as fp:
                data = fp.read()
                data = data.replace("__VENV_NAME__", prompt)
                data = data.replace("__VENV_DIR__", venv_dir)
                data = data.replace("__VENV_BIN_NAME__", venv_bin_name)
                data = data.replace(
                    "__VENV_PYTHON__",
                    os.path.normpath(sys.executable)
                )

            with open(target, "wt") as fp:
                fp.write(data)


def main():
    global verbose
    if '-q' in sys.argv or '-quiet' in sys.argv:  # kept for backward compat
        verbose = False
    if '-v' in sys.argv or '-verbose' in sys.argv or '--verbose' in sys.argv:
        verbose = True

    create_activation_scripts()

    if sys.platform == "win32":
        bin_dir = os.path.join(sys.prefix, 'Scripts')
    else:
        bin_dir = os.path.join(sys.prefix, 'bin')

    if verbose:
        print 'Fixing Python locations in directory: %r' % bin_dir

    for fname in os.listdir(bin_dir):
        path = os.path.join(bin_dir, fname)
        if path.endswith(".exe"):
            try:
                fix_new_launcher(path)
            except LauncherSignatureNotFound:
                if verbose:
                    print("no sig found for {!r}".format(path))
                fix_script(path)
        else:
            fix_script(path)

    if sys.platform != "win32":
        from distutils.sysconfig import get_makefile_filename
        fix_makefile(get_makefile_filename())

        if sys.platform != 'darwin':
            for fn in os.listdir(bin_dir):
                if fn.startswith('remove-'):
                    fix_remove(os.path.join(bin_dir, fn))


if __name__ == '__main__':
    main()