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

The test driver is fully compatible with HPC job submission and all of
ChemShell's usual job submission options can be specified.

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
# TD  29/2/2024 Allow relative paths to be given as --exe option


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

    # YL 18/07/2024: inserted a message of time-consuming RTF initialisation (pickling)
    print(' Note: If the option to install CHARMM data was enabled when ChemShell')
    print('       was compiled, the first time run may take a while to initialise.\n')

    # First check if a filepath has been given
    if os.path.isfile(exe):
        # Ensure absolute path is taken forward
        exe = os.path.abspath(exe)
    # Next check if an executable on the path has been given
    elif shutil.which(exe):
        # No further action required in this case
        pass
    # Finally, search for (chemsh) executable in the standard location 
    else:
        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
    # TD 13/12/2024: Strip trailing "/" as it prevents directory match
    if isinstance(testsets, str):
        testsets = testsets.rstrip('/')
        printsets = testsets
    else:
        testsets = [item.rstrip('/') for item in testsets]
        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
    import os, shutil

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

    list_unsupported = []
    list_supported = []

    # Log of failures in each test directory for output summary
    dir_fails = {}

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

    # TD 21/2/2024: Catch if orca binary in path is the DFT code or the linux screenreader
    # The screenreader is included in many linux distros
    # Annoyingly, the test script will hang indefinitely when orca is called for the first time in this case
    try:
        orca_path = shutil.which("orca")
        orca_size = os.path.getsize(orca_path)
        # The disk sizes of the two binaries are quite different, when I checked:
        # Orca (screenreader) v42.0 :     9272
        # Orca (DFT) v5.0.3         : 52461488
        # Not sure what would cause a return of 0, so exclude it from the check
        if 0 < orca_size < 20000:
            print("\n ORCA in path identified as Linux screen reader -- ORCA tests will be skipped")
            list_unsupported.extend(["orca", "subtractive/orca", "qmmm/orca-gulp", "subtractive/orca-gulp"])
    # If there is no orca in path, shutil.which will return None
    # Then os.path.getsize(None) will raise a TypeError
    except TypeError:
        # No ORCA found in path, tests will fail normally later
        pass

    # 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("/")
        # YL 09/03/2024: there was a bug here that made cases such as 'qmmm/orca-gulp'
        #                ignored when calling `test --sets orca`
        if testsets in [ 'all', ['all'] ] or \
           any(x == testdir for x in testsets) or \
           any(x in 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()

            # YL 24/11/2023: skip if this test is unsupported
            # YL NB 24/11/2023: this is for handling special cases such as in GULP where
            #                   only a single test is unsupported, however this will NOT
            #                   work when the first test is like so; this works for now but
            #                   apparently is not ideal
            # TWK 2/1/2024: keep a separate status and count for skipped tests
            if any([ foo == testdir for foo in list_unsupported ]) and not \
               any([ foo == testdir for foo in list_supported ]):
                rstat, wtime, errmsg = 0, 0.0, ""
                wstat = "skipped"
            else:
                (rstat, wstat, wtime, errmsg) = validate(exelist, testdir, test)

            result = ""
            if wstat == "unsupported":
                # YL 24/11/2023: we remember this type of test is not supported
                if testdir not in list_unsupported:
                    list_unsupported += [ testdir ]
                nunsupported += 1
                result = "M"
                message = "unsupported test (e.g. external code not found)"
            elif wstat == "skipped":
                # TWK 2/1/2024 count skipped tests as a subset of unsupported
                nunsupported += 1
                nskipped += 1
                result = "S"
                message = "skipped test (assumed unsupported)"
            elif (rstat != 0 or wstat == "fail" or wstat == "incomplete"):
                print("-------- FAILED --------- ", end='')
                nfail += 1
                # TWK 4/10/2024 count failures in each directory to summarise later
                dir_fails[testdir] = dir_fails.get(testdir, 0) + 1
                result = "F"
                message = errmsg
            elif wstat == "pass":
                # YL 24/11/2023: we remember this type of test is supported
                if testdir not in list_supported:
                    list_supported += [ testdir ]
                npass += 1
                result = "P"
                message = "OK"
            else:
                nunknown += 1
                result = "U"
                message = "unrecognised result"

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

    # Summary
    summary = f" Failures: {nfail:3}  Passes: {npass:3}  Unsupported: {nunsupported:3}  Unknown: {nunknown:3}"
    skippednote = ""
    if nskipped > 0:
        skippednote = f" Skipped tests (assumed unsupported): {nskipped:3}\n"
    failsummary = ""
    if nfail > 0:
        failsummary = " Summary of failures:\n"
        for faildir in dir_fails:
            failsummary += f" {dir_fails[faildir]:>3} in {faildir}\n"

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

    valfile.write(summary + "\n" + skippednote + "\n" + failsummary)

    valfile.close()

    return (nfail, npass, nunsupported, nunknown)


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

    import subprocess
    from time import time
    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
    runcmd = []
    runcmd.extend(exelist)
    runcmd.append(testname)
    t_start = time()
    runCS = subprocess.run(runcmd, stdout=outfile, stderr=errfile)
    t_end = time()
    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]
    wtime = t_end - t_start
    errfile.close()
    outfile.close()

    # Reset working directory
    chdir(cwd)

    return (rstat, wstat, wtime, errmsg)


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

    from sys  import exit
    from time import time

    t0 = time()

    # 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)

    print(f"\n Time used: {time()-t0:.2f} s\n")

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


if __name__ == '__main__':

    main()

