#!/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.
"""


import mmap
import os.path
import re
import shutil
import stat
import sys
from subprocess import check_output

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(rb"^#!.+$", 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 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 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:
    eups = shutil.which("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)