#!/usr/bin/env python3

#  Copyright (C) 2023 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/>.

"""Py-ChemShell test framework

Simple automated regression test driver suitable for testing both python
and chemsh script execution modes.

As far as possible the test driver runs independently of ChemShell itself,
to replicate standard user execution conditions.

By default all tests in test_suite.txt are run, but a list of directories
to run may optionally be specified.
"""

# TWK 12/7/2023 Converted from original tcsh script to python
# TWK 14/7/2023 Process test set and executable specifications
# TWK 17/7/2023 Add support for chemsh script job submission mode


def processArgs():
    """Process command line options.

Available options:
 --exe    Executable to test (e.g. python3 or chemsh script)
          Default is the chemsh script which will be searched for in the
          usual ../bin/ directory if it is not in the user PATH.
          Optional arguments for the chemsh script can be added and will 
          be passed on. Job submission with chemsh --submit is supported.

 --sets   A list of directories to test.
          Can be an individual directory or subdirectory name, 
          or a full (relative) path.
          Subdirectories will be included, e.g. --sets qmmm will test
          all of the subdirectories of qmmm.
"""

    from argparse import ArgumentParser
    import shutil, os

    parser = ArgumentParser(prog='test')

    # Test set specification
    # A list of directory names
    parser.add_argument('--sets', nargs='*', default='all', 
            help='specify list of test directories to run')

    # Specify executable mode
    # e.g. (path to) python3 or (path to) chemsh script
    parser.add_argument('--exe', nargs='?', default='chemsh', 
            help='set executable to test (e.g. chemsh or python3)')

    # Check if submission to a queueing engine has been requested
    # as this requires special handling
    parser.add_argument('--submit', action='store_true', 
            help='submit tests to queueing system via chemsh script')

    args = parser.parse_known_args() 

    # Get supplied executable
    exe = args[0].exe

    # Search for (chemsh) executable in the standard location if not in path
    if not shutil.which(exe):
        print(" Searching for executable", exe, "...")
        search_path = os.path.abspath("../bin/")
        found = False
        for root, dir, files in os.walk(search_path):
            if exe in files:
                exe = os.path.join(root, exe)
                found = True
                break
        if not found:
            print(" Error: could not find executable")
            exit(1)

    # All unknown arguments are passed on as arguments to the executable
    exelist = args[1]
    exelist.insert(0, exe)
    exestr = ' '.join(exelist)
    print(" Testing executable:", exestr)

    # Make sure PYTHONPATH is set if testing python execution to avoid confusing errors
    if exe.endswith("python") or exe.endswith("python3"):
        if "PYTHONPATH" in os.environ:
            os.environ["PYTHONPATH"] = os.path.abspath("../") + os.pathsep + os.environ["PYTHONPATH"]
        else:
            os.environ["PYTHONPATH"] = os.path.abspath("../")
        print(" Setting PYTHONPATH:", os.environ["PYTHONPATH"])

    # Get requested test sets
    testsets = args[0].sets
    if isinstance(testsets, str):
        printsets = testsets
    else:
        printsets = ' '.join(testsets)
    print("\n Running test sets:", printsets)

    # Check if job submission requested
    submit = args[0].submit

    return (exelist, testsets, submit)


def getTestSuite():
    """Read the list of ChemShell tests from disk into a dictionary."""

    test_suite = open("test_suite.txt", "r")
    testlines = test_suite.read().splitlines()
    test_suite.close()

    # Build a dictionary of the tests
    testdict = {}
    testdir = ""
    dirlist = []
    for line in testlines:
        # Skip any comments
        if line.startswith("#"):
            continue
        elif ":" in line:
            # Save the previous directory list, if it exists
            if testdir != "":
                testdict[testdir] = dirlist
            # Initialise new directory list
            testdir = line[:line.index(":")]
            dirlist = []
        elif ".py" in line:
            # Add test to the current directory list
            dirlist.append(line.strip())
    # Make sure to add the final directory list after end of file
    testdict[testdir] = dirlist

    return testdict


def submitTests(exelist, testsets):
    """Run the test driver as part of a chemsh script job submission."""

    from sys import exit
    import os, subprocess

    print("\n Job submission requested via chemsh script.")
    print(" Please check the output of the submitted job for the test results.")

    # This includes any options passed to chemsh, except for --submit
    runcmd = exelist

    # Pass on the job submission option (which will not be in exelist
    # because it has been parsed).
    runcmd.append("--submit")

    # These options ensure that the test driver is submitted to the queue
    # by the chemsh script, including any specific test set requests.
    runcmd.append("--test")
    runcmd.append(os.path.abspath(__file__))

    runcmd.append("--sets")
    if isinstance(testsets, list):
        runcmd.extend(testsets)
    else:
        runcmd.append(testsets)

    # Launch chemsh and quit (tests will be run via the submitted job).
    # Note that the return code here simply confirms that the chemsh script has
    # been launched and does not relate to the results of the tests themselves.
    runSubmit = subprocess.run(runcmd)
    exit(runSubmit.returncode)


def runTestSuite(exelist, testsets, testdict):
    """Run requested test sets."""

    from sys import stdout

    # Initialise result counters
    nfail = 0
    npass = 0
    nunsupported = 0
    nunknown = 0

    # Initialise validation log file
    valfile = open("validate.log", "w")

    # Loop over the tests in each directory
    for testdir in testdict:
        # Here we check if the directory is one of the requested sets
        # which can be either individual directory names or (relative) paths
        # to directories.
        # So e.g. qmmm, mndo-gulp, and qmmm/mndo-gulp are all valid requests
        runtests = False
        dirnames = testdir.split("/")
        if (testsets == "all" or testsets == ['all'] or
            any(x == testdir for x in testsets) or
            any(x in dirnames for x in testsets)):
                runtests = True
        if runtests == False:
            continue
        print ("\n Test directory:", testdir)
        for test in testdict[testdir]:
            print(f'  - {test} ', end='')
            stdout.flush()

            (rstat, wstat, utime, errmsg) = validate(exelist, testdir, test)

            result = ""
            if wstat == "unsupported":
                nunsupported += 1
                result = "M"
                message = "unsupported test: external code not found"
            elif (rstat != 0 or wstat == "fail" or wstat == "incomplete"):
                print("-------- FAILED --------- ", end='')
                nfail += 1
                result = "F"
                message = errmsg
            elif wstat == "pass":
                npass += 1
                result = "P"
                message = "OK"
            else:
                nunknown += 1
                result = "U"
                message = "unrecognised result"

            print(result, message, utime)
            # Write result to log file
            valfile.write(f'{result}  {testdir:14} {test:24} {utime} {message}\n')

    # Summary
    summary = f" Failures: {nfail:3}  Passes: {npass:3}  Unsupported: {nunsupported:3}  Unknown: {nunknown:3}"

    linewidth = 70
    print("\n "+linewidth*"-")
    print(summary)
    print(" "+linewidth*"-")

    valfile.write(summary + "\n")

    valfile.close()

    return (nfail, npass, nunsupported, nunknown)


def validate(exelist, testdir, testname):
    """Run test and retrieve test status and timing data."""

    import subprocess 
    from pathlib import Path
    from os import chdir

    # Run the test in its own directory
    cwd = Path.cwd()
    testpath = cwd / testdir
    chdir(testpath)
   
    basename = Path(testname).stem
    outfile = open(basename + ".log", "w")
    errfile = open(basename + ".err", "w+")

    # Set to fail unless test explicitly passes
    statusfile = open("test_status", "w")
    statusfile.write("incomplete\n")
    statusfile.close()

    # Run ChemShell (with time measurement)
    runcmd = ['time']
    runcmd.extend(exelist)
    runcmd.append(testname)
    runCS = subprocess.run(runcmd, stdout=outfile, stderr=errfile)
    rstat = runCS.returncode

    # Fetch the result from test_status (written by ChemShell)
    statusfile = open("test_status", "r")
    wstat = statusfile.read().strip()
    statusfile.close()

    # Fetch timing data from output of time command
    # and any error message
    # NB: Use of time command is more reliable than attempting to 
    # measure at the level of this script
    errfile.seek(0)
    errlines = errfile.read().splitlines()
    errmsg = ""
    # Grab a line from error output (before timing data) to print on failure
    # TODO: Maybe we should produce/search for something specific here
    if len(errlines) > 3:
        errmsg = errlines[-4]
    # Chop out the relevant timing data from output of time command
    utime = ""
    for line in errlines:
        # Note we will get the last occurence of user just in case...
        if "user" in line:
            utime = line[:line.index("user")]

    errfile.close()
    outfile.close()

    # Reset working directory
    chdir(cwd)

    return (rstat, wstat, utime, errmsg)


def main():
    """Run the testing framework."""

    from sys import exit

    # Welcome message
    linewidth = 70
    print("\n "+linewidth*"=")
    print(" Py-ChemShell test suite")
    print(" "+linewidth*"=" + "\n")

    # Process command line options
    (exelist, testsets, submit) = processArgs()

    if submit:
        # Run tests via chemsh job submission
        submitTests(exelist, testsets)

    # Import test suite from disk
    testdict = getTestSuite()

    # Run the requested test sets
    (nfail, npass, nunsupported, nunknown) = runTestSuite(exelist, testsets, testdict)

    # Return code for use with continuous integration systems
    if (nfail > 0 or nunknown > 0):
        exit(1)
    else:
        exit(0)


if __name__ == '__main__':

    main()

