#!/usr/bin/env python

""" This script is intended to fixup the shebang paths on scripts installed via
EUPS distrib "tarball" packages.  The assumption is that an end-user will
invoke it as a manual post-install step on OSX systems with SIP installed.
"""

from __future__ import absolute_import, print_function

import io
import mmap
import os.path
import re
import stat
import sys
from subprocess import check_output
from distutils.spawn import find_executable

NON_SCRIPT_EXT = (
    '.pyc', '.pyo', '.h', '.a', '.c', '.cc', '.txt', '.html',
    '.xml', '.png', '.jpg', '.gif', '.class',
    '.o', '.dylib', '.so', '.os', '.version', '.m4', '.table',
    '.pdf', '.tex', '.dox', '.tag', '.conf', '.cfg', '.fits', '.fz',
    '.js', '.yaml', '.md', '.css', '.rst',
)

SHEBANG_PAT = re.compile(br'^#!.+$', re.M)


def fix_shebang(path, build_python):
    if path.endswith(NON_SCRIPT_EXT):
        return

    st = os.stat(path, follow_symlinks=False)
    if stat.S_ISLNK(st.st_mode):
        return
    if not stat.S_ISREG(st.st_mode):
        return
    if st.st_size == 0:
        return
    # must be executable
    if st.st_mode & (stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) == 0:
        return

    # we can't fix unwritable files so they should be completely skipped.
    if not os.access(path, os.R_OK | os.W_OK):
        return

    with io.open(path, mode='rb') as fi:
        data = fi.read(1000)

        if not data.startswith(b'#!'):
            return

        m = SHEBANG_PAT.match(data)
        if not m:
            return

        # skip scripts that use #!/usr/bin/env
        if b'/usr/bin/env' in m.group():
            return
        if b'python' not in m.group():
            return

        fi.seek(0)
        mm = mmap.mmap(fi.fileno(), 0, prot=mmap.PROT_READ)
        data = mm[:]

    encoding = 'utf8'

    # remove the conda prefix logic and set the path to the python interp
    # explicity
    py_exec = (build_python)
    new_data = SHEBANG_PAT.sub(b'#!' + py_exec.encode(encoding), data, count=1)
    if new_data == data:
        return
    print("updating shebang:", path)

    # save original file mode
    mode = os.stat(path).st_mode
    with io.open(path, 'w', encoding=encoding) as fo:
        fo.write(new_data.decode(encoding))
    # restore file mode
    os.chmod(path, mode)


# /post.py

try:
    flavor = os.environ['SHTRON_EUPS_FLAVOR']
except KeyError:
    # py >= 3.3 has shutil.which
    eups = find_executable('eups')
    if not eups:
        raise RuntimeError('unable to find eups command')

    flavor = check_output([eups, "flavor"])
    if not flavor:
        raise RuntimeError('eups flavor sub-command may be broken')

    flavor = flavor.decode('utf-8')

# remove newline from eups output
flavor = flavor.rstrip()

try:
    eups_path = os.environ['EUPS_PATH']
except KeyError:
    raise RuntimeError('required environment variable EUPS_PATH is missing')

prod_install_path = os.path.join(eups_path, flavor)
# sanity check
if not os.path.isdir(prod_install_path):
    raise RuntimeError(f'EUPS product install path "{prod_install_path}"'
                       'is missing')

# fully qualified path to new shebang interp
try:
    py = os.environ['SHTRON_PYTHON']
except KeyError:
    py = sys.executable

for root, dirs, files in os.walk(prod_install_path):
    for f in files:
        path = os.path.join(root, f)
        fix_shebang(path, py)