#!/usr/bin/env python
#
# test_create_remove_wrapper.py - Test the createFSLWrapper.sh and
# removeFSLWrapper.sh scripts. Requires Python>=3.7.
#
# Author: Paul McCarthy <pauldmccarthy@gmail.com>
#


import sys
import os.path as op
import subprocess as sp
import textwrap as tw
import itertools as it
import os
import shlex
import shutil
import tempfile
import contextlib
from unittest import mock


# Paths to create/remove wrapper scripts, when
# running from an in-source version. These may
# be overwritten within the __main__ clause at
# the bottom, where BASE_DIR may be provided
# as a command-line argument.
BASE_DIR       = op.abspath(op.join(op.dirname(__file__), '..'))
CREATE_WRAPPER = op.join(BASE_DIR, 'share', 'fsl', 'sbin', 'createFSLWrapper')
REMOVE_WRAPPER = op.join(BASE_DIR, 'share', 'fsl', 'sbin', 'removeFSLWrapper')


def run(cmd, **kwargs):
    return sp.run(shlex.split(cmd), check=True, **kwargs)


@contextlib.contextmanager
def temp_fsldir():

    testdir    = tempfile.mkdtemp()
    prevdir    = os.getcwd()
    fsldir     = op.join(testdir, 'fsl')
    wrapperdir = op.join(testdir, 'fsl', 'share', 'fsl', 'bin')
    try:

        os.chdir(testdir)
        os.mkdir(fsldir)

        with mock.patch.dict(os.environ, {
                'FSLDIR'                     : fsldir,
                'PREFIX'                     : fsldir,
                'FSL_CREATE_WRAPPER_SCRIPTS' : '1'}):
            yield fsldir, wrapperdir

    finally:
        os.chdir(prevdir)
        shutil.rmtree(testdir)


def touch(path):
    dirname = op.dirname(path)
    if not op.exists(dirname):
        os.makedirs(dirname)
    with open(path, 'wt') as f:
        f.write('.')


def get_called_command(filename):
    """Returns the command that is being called by the given wrapper script.
    """

    with open(filename, 'rt') as f:
        line = f.readlines()[1]

    tokens = line.split()
    cmd    = op.basename(tokens[0])

    if cmd in ('python', 'pythonw'):
        cmd = tokens[2]

    return cmd


def test_env_vars_not_set():
    """Test that wrapper scripts are not created if the
    FSL_CREATE_WRAPPER_SCRIPTS, FSLDIR, or PREFIX environment variables
    are not set.
    """
    with temp_fsldir() as (fsldir, wrapperdir):
        touch(op.join(fsldir, 'bin', 'test_script'))

        env = os.environ.copy()
        env.pop('FSL_CREATE_WRAPPER_SCRIPTS')
        run(f'{CREATE_WRAPPER} test_script', env=env)
        assert not op.exists(op.join(wrapperdir, 'test_script1'))

        env = os.environ.copy()
        env.pop('FSLDIR')
        run(f'{CREATE_WRAPPER} test_script', env=env)
        assert not op.exists(op.join(wrapperdir, 'test_script1'))

        env = os.environ.copy()
        env.pop('PREFIX')
        run(f'{CREATE_WRAPPER} test_script', env=env)
        assert not op.exists(op.join(wrapperdir, 'test_script1'))

        # FSLDIR invalid
        env = os.environ.copy()
        env['FSLDIR'] = '/some/non-existent/path'
        run(f'{CREATE_WRAPPER} test_script', env=env)
        assert not op.exists(op.join(wrapperdir, 'test_script1'))

        # FSLDIR != PREFIX
        env = os.environ.copy()
        env['FSLDIR'] = op.join(env['PREFIX'], 'other_fsl')
        run(f'{CREATE_WRAPPER} test_script', env=env)
        assert not op.exists(op.join(wrapperdir, 'test_script1'))


def test_create_python_wrapper():
    """Test creation of a wrapper script for a python executable"""
    with temp_fsldir() as (fsldir, wrapperdir):

        script_path  = op.join(fsldir, 'bin', 'test_script')
        wrapper_path = op.join(wrapperdir,    'test_script')

        touch(script_path)
        with open(script_path, 'wt') as f:
            f.write('#!/usr/bin/env python\n')
            f.write('print("hello")\n')

        expect = tw.dedent(f"""
        #!/usr/bin/env bash
        {fsldir}/bin/python -I {script_path} "$@"
        """).strip()

        run(f'{CREATE_WRAPPER} test_script')

        assert op.exists(wrapper_path)
        with open(wrapper_path, 'rt') as f:
            got = f.read().strip()

        assert got == expect


def test_create_pythonw_wrapper():
    """Test creation of a wrapper script for a python script which
    uses pythonw as the interpreter (for GUI apps on macOS).
    """
    hashbangs = [
        '#!/bin/bash /usr/bin/pythonw',
        '#!/usr/bin/pythonw',
        '#!/usr/bin/env pythonw']
    with temp_fsldir() as (fsldir, wrapperdir):

        script_path  = op.join(fsldir, 'bin', 'test_script')
        wrapper_path = op.join(wrapperdir,    'test_script')

        touch(script_path)

        for hb in hashbangs:
            with open(script_path, 'wt') as f:
                f.write(hb + '\n')
                f.write('print("hello")\n')

            expect = tw.dedent(f"""
            #!/usr/bin/env bash
            {fsldir}/bin/pythonw -I {script_path} "$@"
            """).strip()

            run(f'{CREATE_WRAPPER} test_script')

            assert op.exists(wrapper_path)
            with open(wrapper_path, 'rt') as f:
                got = f.read().strip()

        assert got == expect


def test_create_other_wrapper():
    """Test creation of a wrapper script for a non-python executable."""
    with temp_fsldir() as (fsldir, wrapperdir):
        script_path  = op.join(fsldir, 'bin', 'test_script')
        wrapper_path = op.join(wrapperdir,    'test_script')

        touch(script_path)
        with open(op.join(fsldir, 'bin', 'test_script'), 'wt') as f:
            f.write('#!/usr/bin/env bash\n')
            f.write('echo "hello"\n')

        expect = tw.dedent(f"""
        #!/usr/bin/env bash
        {script_path} "$@"
        """).strip()

        run(f'{CREATE_WRAPPER} test_script')

        assert op.exists(wrapper_path)
        with open(wrapper_path, 'rt') as f:
            got = f.read().strip()

        assert got == expect



def test_permissions_preserved():
    """Test that wrapper script has same permissions as wrapped script."""
    with temp_fsldir() as (fsldir, wrapperdir):
        perms        = [0o777, 0o755, 0o644, 0o600, 0o755, 0o700]
        script_path  = op.join(fsldir, 'bin', 'test_script')
        wrapper_path = op.join(wrapperdir,    'test_script')
        touch(script_path)

        for perm in perms:
            os.chmod(script_path, perm)
            run(f'{CREATE_WRAPPER} test_script')

            stat = os.stat(wrapper_path)
            assert (stat.st_mode & 0o777) == perm


def test_create_remove_wrappers():
    """Tests normal usage. """
    with temp_fsldir() as (fsldir, wrapperdir):
        touch(op.join(fsldir, 'bin', 'test_script1'))
        touch(op.join(fsldir, 'bin', 'test_script2'))

        run(f'{CREATE_WRAPPER} test_script1 test_script2')

        assert op.exists(op.join(wrapperdir, 'test_script1'))
        assert op.exists(op.join(wrapperdir, 'test_script2'))

        run(f'{REMOVE_WRAPPER} test_script1 test_script2')

        assert not op.exists(op.join(wrapperdir, 'test_script1'))
        assert not op.exists(op.join(wrapperdir, 'test_script2'))


def test_create_gui_wrappers():
    """Tests creation of wrappers for FSL GUI commands, which are called
    "<Command>_gui" on macOS, and "<Command>" on linux, where "<command>"
    (note the case) may also exist. Post-link scripts should only pass the
    "<Command>_gui" variant.
    """

    # Test outcome differs for different platforms.
    # Keys are passed to createFSLWrapper, values are
    # wrappers that should be created
    if sys.platform == 'darwin':
        scripts = {'script'     : 'script',
                   'Script_gui' : 'Script_gui'}

    # linux
    else:
        scripts = {'script'     : 'script',
                   'Script_gui' : 'Script'}

    with temp_fsldir() as (fsldir, wrapperdir):
        for target in scripts.values():
            touch(op.join(fsldir, 'bin', target))

        for wrappers in it.permutations(scripts.keys()):
            args  = ' '.join(wrappers)
            run(f'{CREATE_WRAPPER} {args}')

            for arg in wrappers:
                target  = scripts[arg]
                wrapper = op.join(wrapperdir, target)
                assert op.exists(wrapper)
                assert get_called_command(wrapper) == scripts[arg]

            run(f'{REMOVE_WRAPPER} {args}')
            for arg in wrappers:
                target  = scripts[arg]
                wrapper = op.join(wrapperdir, target)
                assert not op.exists(wrapper)


def test_create_wrappers_no_handle_gui_wrappers():
    """Tests creation of wrappers for FSL commands which might end with "_gui",
    and for which no special handling should take place (see
    test_create_gui_wrappers above).
    """
    scripts = {'fsl'     : 'fsl',
               'fsl_gui' : 'fsl_gui'}

    with temp_fsldir() as (fsldir, wrapperdir):
        for target in scripts.values():
            touch(op.join(fsldir, 'bin', target))

        for wrappers in it.permutations(scripts.keys()):
            args  = ' '.join(wrappers)
            run(f'{CREATE_WRAPPER} {args}')

            for arg in wrappers:
                target  = scripts[arg]
                wrapper = op.join(wrapperdir, target)
                assert op.exists(wrapper)
                assert get_called_command(wrapper) == scripts[arg]

            run(f'{REMOVE_WRAPPER} {args}')
            for arg in wrappers:
                target  = scripts[arg]
                wrapper = op.join(wrapperdir, target)
                assert not op.exists(wrapper)


def test_create_wrappers_rename():
    """Tests the renaming functionality in createFSLWrapper.  If
    $FSLDIR/bin/script exists, a wrapper with a different name
    (e.g. $FSLDIR/share/fsl/bin/renamed_script) can be created by passing
    "script=renamed_script".
    """

    # Keys are passed to createFSLWrapper, values
    # are wrappers that should be created
    scripts = {
        'script1=renamed_script1'         : 'renamed_script1',
        'script2=renamed_script2'         : 'renamed_script2',
        'script3_gui=renamed_script3_gui' : 'renamed_script3_gui',
        'script4_gui=renamed_script4'     : 'renamed_script4'
    }

    with temp_fsldir() as (fsldir, wrapperdir):
        for script in scripts.keys():
            target = script.split('=')[0]
            with open(target, 'wt') as f:
                touch(op.join(fsldir, 'bin', target))

        for wrappers in it.permutations(scripts.keys()):
            args  = ' '.join(wrappers)
            run(f'{CREATE_WRAPPER} {args}')

            for arg in wrappers:
                target  = arg.split('=')[0]
                wrapper = op.join(wrapperdir, scripts[arg])

                assert op.exists(wrapper)
                assert get_called_command(wrapper) == target

            run(f'{REMOVE_WRAPPER} {args}')
            for arg in wrappers:
                target  = scripts[arg]
                wrapper = op.join(wrapperdir, target)
                assert not op.exists(wrapper)


def test_create_wrappers_multiple_same():
    """Tests creating multiple wrapper scripts which call the same
    target command.
    """

    # Keys are passed to createFSLWrapper, values
    # are wrappers that should be created
    scripts = {
        'scripta'         : 'scripta',
        'scripta=script1' : 'script1',
        'scripta=script2' : 'script2',
        'scriptb'         : 'scriptb',
        'scriptc=script3' : 'script3',
        'scriptc=script4' : 'script4',
    }

    with temp_fsldir() as (fsldir, wrapperdir):
        for script in scripts.keys():
            target = script.split('=')[0]
            with open(target, 'wt') as f:
                touch(op.join(fsldir, 'bin', target))

        for wrappers in it.permutations(scripts.keys()):
            args  = ' '.join(wrappers)
            run(f'{CREATE_WRAPPER} {args}')

            for arg in wrappers:
                target  = arg.split('=')[0]
                wrapper = op.join(wrapperdir, scripts[arg])

                assert op.exists(wrapper)
                assert get_called_command(wrapper) == target

            run(f'{REMOVE_WRAPPER} {args}')
            for arg in wrappers:
                target  = scripts[arg]
                wrapper = op.join(wrapperdir, target)
                assert not op.exists(wrapper)


if __name__ == '__main__':
    # base dir can be speecified on command line
    if len(sys.argv) > 1:
        BASE_DIR       = sys.argv[1]
        SBIN_DIR       = op.join(BASE_DIR, 'share', 'fsl', 'sbin')
        CREATE_WRAPPER = op.join(SBIN_DIR, 'createFSLWrapper')
        REMOVE_WRAPPER = op.join(SBIN_DIR, 'removeFSLWrapper')

    thismod = sys.modules[__name__]
    for testname in dir(thismod):
        if testname.startswith('test_'):
            test = getattr(thismod, testname)
            if callable(test):
                print(f'Running test {testname}')
                test()