#  Copyright (C) 2017 The authors of Py-ChemShell
#
#  This file is part of Py-ChemShell.
#
#  Py-ChemShell is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation, either version 3 of the
#  License, or (at your option) any later version.
#
#  Py-ChemShell is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with Py-ChemShell.  If not, see
#  <http://www.gnu.org/licenses/>.

# YL 17/02/2019: to generate a "here document" chemsh script at compiletime: no blank line should be left after this header 
#                NB: have to make it a list because not all versions of CMake like newlines in a string
#                NB: `python3 << - $*` won't allow an interactive Python shell
__bash_footer = [ 'EOF', '\\\n', 'python3 -c \\\"\$CMD\\\" \$*' ]
# moved to the specs files as '{}'.format() required: __bash_header = [ '#!/bin/bash --login', 'module load {}', 'read -r -d \"\" CMD << EOF' ]

def init(args):
    '''Initialise platform environment'''

    from os.path import join

    # reference to a feature setter function
    args.__setarg = __setarg
    __setarg.__args = args

    # load platform's specification module
    specsmodule = args.__import(args.platform, join(args.__CHEMSH_ROOT, 'chemsh', 'utils', 'platforms', args.platform+'.py'))

    # initialise
    specsmodule.init(args)


def setup(args):
    '''Prepare platform for setup'''

    from os.path import join

    # load platform's specification modul
    specsmodule = args.__import(args.platform, join(args.__CHEMSH_ROOT, 'chemsh', 'utils', 'platforms', args.platform+'.py'))

    # initialise
    specsmodule.setup(args)


def challengeSerial(args):
    ''''''

    try:
        if not args.mpi:
            print(" >>> WARNING: serial ChemShell (built without --mpi argument) may not be able to run on parallel compute nodes\n")
            promptYes()
    except:
        pass


def promptYes():
    '''Press y or Enter to continue'''

    keybuff = input("\n Are you sure to continue? (Y/n)\n")

    if keybuff.lower() in [ '', 'y', 'yes' ]:
        pass
    else:
        print("\n >>> Exiting ChemShell ARCHER Environment...\n")
        # FIXME: doesn't quit
        exit(0)


def getEnvModules(moduleshome):
    '''Environment Modules system'''

    from os             import environ, path, sep
    from importlib.util import spec_from_file_location, module_from_spec

    # Environment Modules
    environ.update(MODULESHOME=moduleshome)
    realpath = path.realpath(environ['MODULESHOME'])
    environ.update(MODULE_VERSION=realpath.rpartition(sep)[-1])
    environ.update(MODULE_VERSION_STACK=environ['MODULE_VERSION'])

    if not environ.get('MODULEPATH'):
        listbuff = []
        with open(path.join(environ['MODULESHOME'], 'init', '.modulespath'), 'r') as fp:
            for line in fp.readlines():
                line = sub("#.*$", '', line)
                if line:
                    listbuff.append(line)
        # FIXME: this it not cross-platform!
        environ['MODULEPATH'] = ':'.join(listbuff)

    environ.setdefault('LOADEDMODULES', '')
#    environ.setdefault('ARMCI_DEFAULT_SHMMAX_UBOUND', '131072')

    # import module.py (from . import won't work because parent module may have not been imported)
    # path.realpath(__file__) crashes on Neale!
    spec = spec_from_file_location('module', path.join(path.dirname(__file__), 'module.py'))
    module = module_from_spec(spec)
    spec.loader.exec_module(module)

    return module


# YL NB 28/01/2021: we can't use NumPy yet!
def getModulesToLoad(defaults, toloads=[]):
    ''''''

    lbuff = []

    # get modules to load
    for key, val in defaults.items():
        # do not add if overridden by user
        add = True
        for m in toloads:
            if m.startswith(key):
                add = False
                break
        if add:
            lbuff += [val]

    toloads += lbuff

    return toloads


def __setarg(key, val):
    '''This only works when not overridden by the user'''

    # __args and __get_default should be defined externally
    try:
        # if still the default value
        if key in __setarg.__args and getattr(__setarg.__args, key) == __setarg.__args.__get_default(key):
            setattr(__setarg.__args, key, val)
    except:
        pass


# this run doesn't invoke chemsh.x
def runPlot(args):
    '''Plot graphs'''

    from subprocess import run
    from os         import environ, path

    if args.mpiexec:
        mpiexec = ' '.join([args.mpiexec, '-n 1'])
    else:
        mpiexec = ''

    # YL 22/04/2022: we have to use the environment variable because `chemsh` is not initialised until
    #                the plotting functions import other ChemShell modules but within `chemsh`
    #                CHEMSH_ARCH will be set to "gnu" as default so that scanning of package modules fails
    dir_exec = path.split(args.__CHEMSH_EXEC)[0]
    if path.isdir(dir_exec):
        CHEMSH_ARCH = path.split(dir_exec)[1]
        environ['CHEMSH_ARCH'] = CHEMSH_ARCH

    plotmodule = args.__import('chemsh.tools.plot', path.join(args.__CHEMSH_ROOT, 'chemsh', 'tools', 'plot.py'), debug=args.debug)

    plotmodule.run(args)

    print("\n ChemShell exiting normally.\n")

    return 0


def runSerial(args):
    '''Run a serial job'''

    from os         import environ, path
    from os.path    import isfile
    from subprocess import run, CompletedProcess, CalledProcessError, STDOUT
    from sys        import exit, stdout, stderr

    # check if chemsh.x exists
    if not isfile(args.__CHEMSH_EXEC):
        print("\n >>> ChemShell ERROR:", args.__CHEMSH_EXEC, "does not exist. It is likely ChemShell was not properly compiled.\n")
        exit(999)

    # YL 10/03/2021: we return an error code for every input file; it might not make sense but the returned ierrors are not in use, either
    ierrors = []

    # a serial job may also be launched by mpirun
    if args.mpiexec:
        mpiexec = ' '.join([args.mpiexec, '-n 1'])
    else:
        mpiexec = ''
    
    # YL 19/09/2019: supported a sequence of multiple jobs
    for filename in args.filenames:

        # interactive mode is triggered if filename=='' (namely args.filenames=[''])
        if filename and not isfile(filename):
            print("\n >>> ChemShell ERROR:", filename, "does not exist. Please check your command.\n")
            ierrors.append(999)
            continue

        try:
            if str(args.prepend_pythonpath):
                # YL 01/03/2021: we reverse the order to make sure the elements are inserted in order (by initpy.c)
                prepend_pythonpath = ':'.join(str(args.prepend_pythonpath).split(':')[::-1])
                lbuff = [ args.wrapper,
                          # mpiexec is empty if built serial
                          mpiexec,
                          args.__CHEMSH_EXEC,
                         '-dbg', str(args.debug),
                         '--prepend-pythonpath', prepend_pythonpath,
                          filename ]
            else:
                lbuff = [ args.wrapper,
                          # mpiexec is empty if built serial
                          mpiexec,
                          args.__CHEMSH_EXEC,
                         '-dbg', str(args.debug),
                          filename ]
            command = ' '.join(lbuff)
            print("\n >>> Executing command:\n {}\n".format(command.strip()))
            stdout.flush()
    
            completedproc = run(command,
                                shell=True,
                                check=True)

# YL 10/03/2021: moved to chemsh.x because ChemShell is not always invoked by chemsh
#                # Message to mark a normal termination
#                print(" ChemShell exiting normally.")


            stdout.flush()
            stderr.flush()

            # YL 10/03/2021: do not return immediately
            ierrors.append(completedproc.returncode)
    
        except CalledProcessError as err:
            if err.returncode:
                print(" >>> Serial run failed. Exiting...\n\n Exit code:", err.returncode)
                ierrors.append(err.returncode)

    return ierrors


def runParallel(args):
    '''Run a parallel job'''

    from os         import environ
    from os.path    import isfile, join
    from shlex      import split as shlex_split
    from subprocess import run, CompletedProcess, CalledProcessError, STDOUT

    # YL 10/03/2021: we return an error code for every input file; it might not make sense but the returned ierrors are not in use, either
    ierrors = []

    # check if chemsh.x exists
    if not isfile(args.__CHEMSH_EXEC):
        print("\n >>> ChemShell ERROR:", args.__CHEMSH_EXEC, "does not exist. It is likely ChemShell was not properly compiled.\n")
        exit(999)

    # YL 19/09/2019: supported a sequence of multiple jobs
    for filename in args.filenames:

        # interactive mode is triggered if filename=='' (namely args.filenames=[''])
        if filename and not isfile(filename):
            print("\n >>> ChemShell ERROR:", filename, "does not exist. Please check your command.\n")
            ierrors.append(999)
            continue

        try:
            if str(args.prepend_pythonpath):
                # YL 01/03/2021: we reverse the order to make sure the elements are inserted in order (by initpy.c)
                prepend_pythonpath = ':'.join(str(args.prepend_pythonpath).split(':')[::-1])
                lbuff = [ args.wrapper,
                          args.mpiexec,
                         '-n'   , str(args.nprocs),
                          args.__CHEMSH_EXEC,
                         '-nwg' , str(args.nworkgroups),
                         '-npmm', str(args.nprocs_mm),
                         '-gm'  , str(args.ga_memory),
                         '-dbg' , str(args.debug),
                         '--prepend-pythonpath', prepend_pythonpath,
                          filename ]
            else:
                lbuff = [ args.wrapper,
                          args.mpiexec,
                         '-n'   , str(args.nprocs),
                          args.__CHEMSH_EXEC,
                         '-nwg' , str(args.nworkgroups),
                         '-npmm', str(args.nprocs_mm),
                         '-gm'  , str(args.ga_memory),
                         '-dbg' , str(args.debug),
                          filename ]
            command = ' '.join(lbuff)
            print(" >>> Executing command: {}\n".format(command.strip()))
    
            completedproc = run(command,
                                shell=True,
                                check=True)

# YL 10/03/2021: moved to chemsh.x because ChemShell is not always invoked by chemsh
#                # Message to mark a normal termination
#                print(" ChemShell exiting normally.")
            ierrors.append(completedproc.returncode)
    
        except CalledProcessError as err:
            if err.returncode:
                print(" >>> Parallel run failed. Exiting...\n")
                ierrors.append(err.returncode)


def run(args):
    '''Run ChemShell'''

    from os      import environ
    from os.path import join

    # OpenMP threads
    environ.setdefault('OMP_NUM_THREADS', f'{args.nthreads}')

    # load platform's specification modul
    specsmodule = args.__import(args.platform, join(args.__CHEMSH_ROOT, 'chemsh', 'utils', 'platforms', args.platform+'.py'))

    # a plotting job only
    if args.plot:

        return runPlot(args)

    # for a PC (such as Ubuntu) args.submit is always False
    if args.submit is True:

        # run chemshell queue submitter
        # YL 10/03/2021: i haven't thought of how to obtain all the error codes of a series of jobs submitted
        #                but ierrors now are not used anyway
        ierrors = [specsmodule.run(args)]

    else:

        nprocs = max(args.nprocs, args.nworkgroups)
    
        # if np=0 or np=1 run in serial; if np>1 run in parallel
        selectcases = { 0:runSerial,
                        1:runSerial }
    
        # force runSerial if built in serial (mpiexec is '')
        if not args.mpiexec:
            nprocs = 1
    
        # YL 10/03/2021: now return a list of ierrors, which are not actually used though
        ierrors = selectcases.get(nprocs, runParallel)(args)

    return ierrors


def textwrap_indent(text, indent=' ', width=70):
    '''Wrap a text'''

    from textwrap import wrap
    from textwrap import indent as textwrap_indent

    lindent = len(indent)

    return textwrap_indent('\n'.join(wrap(text, width=width-lindent)), indent)
