#!/usr/bin/env python3
import os
import re
import sys
from optparse import OptionParser

SCRIPT_NAME = os.path.basename(sys.argv[0])
SCRIPT_DIR = os.path.dirname(sys.argv[0])
SCRIPT_PARENT_DIR = os.path.realpath(os.path.join(SCRIPT_DIR, '..'))


ANSI_CURSOR_TO_COLUMN_60 = "\x1b[999D\x1b[60C"
ANSI_COLOR_BRIGHT_GREEN = "\x1b[32;1m"
ANSI_COLOR_BRIGHT_RED = "\x1b[31;1m"
ANSI_COLOR_NORMAL = "\x1b[0;m"


def print_nonl(*args):
    sys.stdout.write(" ".join(args))
    sys.stdout.flush()


def success():
    if sys.stdout.isatty():
        print(
            ANSI_CURSOR_TO_COLUMN_60
            + "["
            + ANSI_COLOR_BRIGHT_GREEN
            + " OK "
            + ANSI_COLOR_NORMAL
            + "]"
        )
    else:
        print("[ OK ]")


def failure(message=""):
    if sys.stdout.isatty():
        print(
            ANSI_CURSOR_TO_COLUMN_60
            + "["
            + ANSI_COLOR_BRIGHT_RED
            + "FAIL"
            + ANSI_COLOR_NORMAL
            + "]"
        )
    else:
        print("[FAIL]")
    if message:
        print(message)


def osg_files_dir(staging_dir):
    return os.path.join(staging_dir, 'portable-xrootd')


def is_valid_staging_dir(staging_dir):
    ofd = osg_files_dir(staging_dir)
    for osg_file in ['setup.sh.in', 'setup.csh.in', 'tarball-run.in']:
        if not os.path.exists(os.path.join(ofd, osg_file)):
            return False
    return True


def write_setup_from_templates(staging_dir, final_osg_location):
    abs_staging_dir = os.path.abspath(staging_dir)

    print("Creating environment setup files...")
    for setup_file, mode in (
        ('setup.sh', 0o644),
        ('setup.csh', 0o644),
        ('tarball-run', 0o755),
    ):
        setup_in_file = setup_file + ".in"
        setup_in_path = os.path.join(abs_staging_dir, "portable-xrootd", setup_in_file)
        setup_path = os.path.join(abs_staging_dir, setup_file)

        if not os.path.exists(setup_in_path):
            failure("%r not found" % (setup_in_path))
            return

        print_nonl("Creating %r" % setup_file)
        setup_fh = None
        setup_in_fh = None
        try:
            setup_fh = open(setup_path, 'w')
            setup_in_fh = open(setup_in_path, 'r')
            if (
                mode != 0o755
            ):  # Executables have shebang lines so don't prepend the comment
                setup_fh.write(
                    """\
# This file was automatically generated from %s by %s
# Rerunning %s will cause modifications to be lost.
"""
                    % (setup_in_path, SCRIPT_NAME, SCRIPT_NAME)
                )
            for in_line in setup_in_fh:
                setup_fh.write(
                    re.sub(r'@@XROOTD_LOCATION@@', final_osg_location, in_line)
                )
        except EnvironmentError as err:
            failure(
                "Unable to write environment setup file for the following reason:\n%s"
                % err
            )
        finally:
            if setup_fh:
                setup_fh.close()
            if setup_in_fh:
                setup_in_fh.close()
        os.chmod(setup_path, mode)
        success()
    # end for


def write_setup_local_files(staging_dir):
    abs_staging_dir = os.path.abspath(staging_dir)

    for shell in 'sh', 'csh':
        setup_local_file = 'setup-local.' + shell
        setup_local_path = os.path.join(abs_staging_dir, setup_local_file)
        if not os.path.exists(setup_local_path):
            setup_local_fh = None
            try:
                print_nonl("Creating %r" % setup_local_file)
                setup_local_fh = open(setup_local_path, 'w')
                setup_local_fh.write(
                    """\
# This file is for local environment customizations. It is sourced at the end
# of setup.%s and will not be overwritten by future runs of %s.
"""
                    % (shell, SCRIPT_NAME)
                )
                setup_local_fh.close()
            except EnvironmentError as err:
                failure(
                    "Unable to write local environment setup file for the following reason:\n%s"
                    % err
                )
            finally:
                if setup_local_fh:
                    setup_local_fh.close()
            success()


def parse_cmdline_args(argv):
    parser = OptionParser(
        """
    %%prog [<STAGING_DIR>] [--final-location=<DIR>]

If STAGING_DIR is not specified on the command line, then the parent
directory of this script (%r) is used for STAGING_DIR.

This script runs the post-extraction steps necessary to complete your install
of the tarball client. This consists of creating environment files (such as
setup.sh).

By default the script assumes that the directory the tarball was extracted to
is the directory the installation will run from. (That is, XROOTD_LOCATION is the
same as STAGING_DIR). If this is not the case, for example if you're extracting
the tarball into a staging area before pushing it out to a network share, then
you must specify the --final-location argument.
"""
        % (SCRIPT_PARENT_DIR)
    )

    parser.add_option(
        "-f",
        "--final-location",
        default=None,
        help="The final location that the software will "
        "be run from. If not specified, the staging dir will be used.",
    )

    options, args = parser.parse_args(argv[1:])

    return (options, args)


def get_staging_dir(arg_staging_dir=None):
    staging_dir = None

    if arg_staging_dir:
        staging_dir_candidate = os.path.abspath(arg_staging_dir)
        print_nonl("Staging dir specified as %r... " % (staging_dir_candidate))
        if is_valid_staging_dir(staging_dir_candidate):
            print("ok")
            staging_dir = staging_dir_candidate
        else:
            print("but is not valid.")
    else:
        print_nonl("Staging dir not specified; trying %r..." % (SCRIPT_PARENT_DIR))
        if is_valid_staging_dir(SCRIPT_PARENT_DIR):
            print("ok")
            staging_dir = SCRIPT_PARENT_DIR
        else:
            print("not valid")

    return staging_dir


def main(argv):
    options, args = parse_cmdline_args(argv)

    if len(args) > 0:
        staging_dir = get_staging_dir(args[0])
    else:
        staging_dir = get_staging_dir(None)

    if not staging_dir:
        print("No valid staging directory found.")
        return 2

    if options.final_location:
        print("Final XROOTD_LOCATION specified as %r" % (options.final_osg_location))
        final_location = options.final_osg_location
    else:
        print(
            "Final XROOTD_LOCATION not specified. Using staging dir (%r)."
            % (staging_dir)
        )
        final_location = staging_dir

    write_setup_from_templates(staging_dir, final_location)
    write_setup_local_files(staging_dir)
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))
