#!/usr/bin/env python2
from __future__ import print_function
import glob
import re
import os
import shutil
import sys
import tempfile

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, 'osg')


def is_valid_staging_dir(staging_dir):
    ofd = osg_files_dir(staging_dir)
    for osg_file in ['setup.sh.in', 'setup.csh.in', 'osgrun.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), ('osgrun', 0o755):
        setup_in_file = setup_file + ".in"
        setup_in_path = os.path.join(abs_staging_dir, "osg", 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'@@OSG_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 fix_osg_location_in_file(file_path, osg_location):
    """Given a path to a file, replace the string @@OSG_LOCATION@@ in the
    file with the path to the directory the software will be run from (the
    final OSG_LOCATION).
    This is basically a sed one-liner written in 20 lines of Python.

    """
    tmp_fh = tmp_path = file_fh = None
    try:
        # mkstemp returns a file descriptor instead of a file object
        _tmp_fd, tmp_path = tempfile.mkstemp()
        tmp_fh = os.fdopen(_tmp_fd, 'w')
        file_mode = os.stat(file_path).st_mode
        file_fh = open(file_path, 'r')
        for file_line in file_fh:
            tmp_fh.write(re.sub(r'@@OSG_LOCATION@@', osg_location, file_line))
        tmp_fh.close()
        if os.path.exists(file_path):
            shutil.copy(file_path, file_path + ".bak")
        shutil.move(tmp_path, file_path)
        os.chmod(file_path, file_mode)
    finally:
        if tmp_fh:
            tmp_fh.close()
        if os.path.exists(tmp_path):
            os.remove(tmp_path)
        if file_fh:
            file_fh.close()


def fix_osg_location_in_fetch_crl(staging_dir, final_osg_location):
    """Write the final OSG_LOCATION directory into the fetch-crl config
    files under the staging dir.

    """
    abs_staging_dir = os.path.abspath(staging_dir)
    etc = os.path.join(abs_staging_dir, 'etc')

    confs = glob.glob(os.path.join(etc, "fetch-crl.conf"))
    if not confs:
        return

    print_nonl("Updating fetch-crl config file(s)")
    # dereference symlinks
    confs_real = set([os.path.realpath(x) for x in confs])
    try:
        # fix each file only once
        for in_path in confs_real:
            fix_osg_location_in_file(in_path, final_osg_location)
        success()
    except EnvironmentError as err:
        failure("Unable to fix fetch-crl config file(s) for the following reason:\n%s" % err)


def fix_osg_location_in_sysconfig_bestman2(staging_dir, final_osg_location):
    """Write the final OSG_LOCATION directory into the sysconfig file
    for bestman2 under the staging dir.

    """
    sysconfig_path = os.path.join(staging_dir, 'etc/sysconfig/bestman2')
    if not os.path.exists(sysconfig_path):
        return

    print_nonl("Updating BeSTMan2 sysconfig file")

    try:
        fix_osg_location_in_file(sysconfig_path, final_osg_location)
        success()
    except EnvironmentError as err:
        failure("Unable to fix BeSTMan2 sysconfig file for the following reason:\n%s" % err)


def parse_cmdline_args(argv):
    parser = OptionParser("""
    %%prog [<STAGING_DIR>] [--final-osg-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) and fixing the installation directory in configuration files (such
as etc/fetch-crl.conf).

By default the script assumes that the directory the tarball was extracted to
is the directory the installation will run from. (That is, OSG_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-osg-location argument.
""" % (SCRIPT_PARENT_DIR))

    parser.add_option("-f", "--final-osg-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_osg_location:
        print("Final OSG_LOCATION specified as %r" % (options.final_osg_location))
        final_osg_location = options.final_osg_location
    else:
        print("Final OSG_LOCATION not specified. Using staging dir (%r)." % (staging_dir))
        final_osg_location = staging_dir

    write_setup_from_templates(staging_dir, final_osg_location)
    write_setup_local_files(staging_dir)
    fix_osg_location_in_fetch_crl(staging_dir, final_osg_location)
    fix_osg_location_in_sysconfig_bestman2(staging_dir, final_osg_location)
    return 0

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

