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


import chemsh
from ..dl_py2f       import utils
from ..dl_py2f.utils import miscutils
from ..utils         import testutils


# FIXME: add returncode
#@miscutils.time
def runExec(theory, syscmd, options=None, dryrun=False, mute=False, stderr='', popen=False, shell=False, **kwargs):
    '''Run an external executable'''

    from sys        import stderr, stdout
    from subprocess import Popen, run
    from .          import errors

    # YL 04/05/2023: added a user-controllable keyword _sys_call_popen at the Interface
    #                level to allow a system call to be executed with a Popen
    #                instance
    # YL 07/06/2023: added _sys_call_shell similarly
    # TWK 15/5/2024: these settings, which override the corresponding arguments to runExec,
    # appear to break e.g. parallel ORCA, which tries to set shell to True via the call.
    # As a workaround have added a condition to set these only if popen/shell are not already
    # set to a non-default value, but probably needs further thought.
    if popen == False:
        popen = theory._sys_call_popen
    if shell == False:
        shell = theory._sys_call_shell

    # JCRT Allow for no arguments 
    if theory.runargs == None:
        theory.runargs = ""

    if chemsh.CHEMSH_DEBUG > 4:
        print('\n >>> ChemShell runExec():')
        print(' theory  = {}'.format(theory))
        print(' syscmd  = {}'.format(syscmd))
        print(' runargs = {}'.format(theory.runargs))

    # assign filename of output
    # AJL, March 2019: This command overwrites any existing defintion of output.
    #                  This doesn't work when the output name is different from the input.
    #                  e.g. FHI-aims. There is no easy fix, but I'm going to put an if statement
    # JCRT: added this to allow for the inputs == None
    if theory.input != None:
        if not theory.output:
            theory.output = utils.fileutils.getOutName(theory.input)

        # generate input
        theory.genInput(**kwargs)

    # a dry run only generates the input(s) but does not actually invoke the calculation
    if dryrun:
        print("\n >>> A dry run of", type(theory).__name__, "performed (only generated input files).")
        return 0

    # run binary in frontground
    with open(theory.output, 'w') as fp_out:

        # error logger
        try:
            fp_err = open(stderr, 'w')
        except:
            fp_err = stderr

        # AJL, March 2019: Made this just a little informative for user by adding output
        #      Oct, 2019 : DL_POLY has theory.runargs as a tuple, not a str, so need to accommodate this.
        if type(theory.runargs) is tuple:
            cmd_str = theory.sysmpi + ' ' + syscmd + ' > ' + theory.output
        elif theory.runargs == "<":
            cmd_str = theory.sysmpi + ' ' + syscmd + ' < ' + theory.input + ' > ' + theory.output
        else:
            cmd_str = theory.sysmpi + ' ' + syscmd + ' ' + theory.runargs + ' > ' + theory.output
        cmd_str = cmd_str.strip()
        print(f' >>> Executing system call command for {type(theory).__name__}:\n {cmd_str}')
        stdout.flush()

        # TWK: Use runargs as invocation may not be as simple as the input file name
        try:

            # YL 11/06/2019: some system commands (e.g., DL_MONTE.X) don't run with communicate(). "input=" not properly defined?
            # YL 21/04/2020: changed to run() to capture error code
            # YL 09/11/2021: refactored to enable shell=True for special uses such as ORCA to pass in quotes
            # TWK 05/08/2020 Handle a request to redirect stdin (requested by setting runargs as "<")
            if type(theory.runargs) is tuple or theory.runargs == '<':
                # plain string if shell=True
                if shell:
                    cmd_buff = theory.sysmpi + ' ' + syscmd
                # list invocation if shell=False
                else:
                    cmd_buff = (theory.sysmpi + ' ' + syscmd).split()
            else:
                # plain string if shell=True
                if shell:
                    cmd_buff = theory.sysmpi + ' ' + syscmd + ' ' + theory.runargs
                # list invocation if shell=False
                else:
                    cmd_buff = (theory.sysmpi + ' ' + syscmd + ' ' + theory.runargs).split()
            # file pointer
            if theory.runargs == '<':
                fp_inp = open(theory.input)
            else:
                fp_inp = None
            # run
            if popen:
                popen = Popen(cmd_buff, stdin=fp_inp, stdout=fp_out, stderr=fp_err, shell=shell)
            else:
                result = run(cmd_buff, stdin=fp_inp, stdout=fp_out, stderr=fp_err, shell=shell)
            # file pointer
            if theory.runargs == '<':
                fp_inp.close()

            # YL 10/11/2020: capture the error output since some binaries issue a `stop` command so that return code is 0!
            #            NB: run(stderr=xxx) has to be a real file (StringIO() or BytesIO() won't be accepted)
            # treat per interface
            try:
                ierror_logger = theory.processErr(stderr)
            except:
                ierror_logger = 0

            # wait for process to terminate 
            if popen:
                outs, ierror = popen.communicate()
            else:
                ierror = ierror_logger or result.returncode

            # close before exiting
            if fp_err != stderr:
                fp_err.close()
            stdout.flush()

            if not ierror:
                print("\n >>>", type(theory).__name__, "ran successfully.\n")
            else:
                # shell not found code
                if ierror == errors.UNSUPPORTED:
                    print(f'\n >>> ERROR: File not found! Failed running {type(theory).__name__} with command\n {cmd_str}\n')
                    # FIXME: should be used with a parameter test=True for test mode when the Python testing system is in place
                    testutils.setUnsupported()
                    exit(errors.UNSUPPORTED)
                else:
                    print(f'\n >>> ERROR: Failed running {type(theory).__name__} with command\n {cmd_str}\n')
                    exit(ierror)

        # these exceptions is for shell=False above
        except (FileNotFoundError, PermissionError):
            print(f'\n >>> ERROR: File not found! Failed running {type(theory).__name__} with command\n {cmd_str}\n')
            # FIXME: should be used with a parameter test=True for test mode when the Python testing system is in place
            testutils.setUnsupported()
            exit(errors.UNSUPPORTED)

    # analyse results
    theory.parseOutput()
    stdout.flush()

    # JCRT Allow for no arguments 
    if theory.runargs == "":
        theory.runargs = None

    return errors.PASSED


#@miscutils.time
def runLib(theory, *args, options=None, dryrun=False, **kwargs):
    '''Run with an internal shared library'''

    # in case a library is called with inp/out (e.g., GULP)
    # JCRT added the second check to allow for programs that have input == None
    if hasattr(theory, 'genInput') and (theory.input != None):
        theory.genInput(**kwargs)

    # a dry run only generates the input(s) but does not actually invoke the calculation
    if dryrun:
        print("\n >>> A dry run of", type(theory).__name__, "performed (only generated input files).")
        return 0

    ierror = theory.runLib(options=options, **kwargs)

    # in case a library is called with inp/out (e.g., GULP)
    if hasattr(theory, 'parseOutput'):
        theory.parseOutput()

    if ierror:
        return 1
    else:
        if not kwargs.get('mute', False):
            print("\n >>>", type(theory).__name__, "run successfully.")
        return 0


@miscutils.memusage
def run(theory, cmdpath, options=None, **kwargs):

#    from resource import getrusage, RUSAGE_SELF
    from ctypes import c_long, CDLL
    from types  import ModuleType
    from ..io   import keyword

    from copy import deepcopy
    '''Run calculation'''

    # JCRT: Multiple executables as a tuple 
    # we can just call this routine recursively
    if type(cmdpath) == tuple:
        if type(theory.output) != tuple and type(theory.input) != tuple and len(theory.output) != len(cmdpath) and len(theory.input) != len(cmdpath):
            raise IOError("cmdpath and theory.output must both be either type(str) or type(tuple) of equal lengths")

        # keep the original member variables
        theory_output = deepcopy(theory.output)
        theory_input = deepcopy(theory.input)

        for i, cmd in enumerate(cmdpath):
            # replace the member variables (tuples) -> strings
            theory.input = theory_input[i]
            theory.output = theory_output[i]

            # run each of the executables with the desired inputs/outputs
            run(theory, cmd, options=options, **kwargs)

        # reset the inputs theory inputs and outputs to the original values
        theory.output = theory_output
        theory.input = theory_input

        # return here to avoid processing the cmdpath directly as a tuple
        return c_long(0)

    selectcases = {
                    str       : runExec,
                    ModuleType: runLib,
                    CDLL      : runLib
                  }

    if chemsh.CHEMSH_DEBUG > 3:
        print("\n\n >>>", theory.__class__.__name__, "fragment (a.u.):\n")
        print(theory.frag.coords2str())
    if chemsh.CHEMSH_DEBUG > 3 and theory.frag.nbqs:
        print("\n\n >>>", theory.__class__.__name__, "background point charges (a.u.):\n")
        print(theory.frag.bqs.coords2str(atoms=None, indent=" "*6))

    # all keyword arguments at the theory.run() step will be temporary
    # NOTE: do not use `__dict__` but only dir(), as all attibutes of <property> are not contained in the latter
    for key in set(kwargs) & set(dir(theory)):
        # swap values (syntax a,b=c,d does not apply)
        buffer = getattr(theory, key)
        setattr(theory, key, kwargs[key])
        kwargs[key] = buffer

    # higher level task information should not be specified in low level interface object
    # runLib or runExec
    ierror = selectcases[type(cmdpath)](theory, cmdpath, options=options, **kwargs)

    for key in set(kwargs) & set(dir(theory)):
        # swap values back (syntax a,b=c,d does not apply)
        buffer = getattr(theory, key)
        setattr(theory, key, kwargs[key])
        kwargs[key] = buffer

#    print("\n >>> Peak memory used: %12.3f MB\n"%(getrusage(RUSAGE_SELF).ru_maxrss/1024/1024))

    # return 0 of ctypes if successful, for Fortran callback
    return c_long(0)
