# Generate CMake test from pytests
#
# Usage:
# python pytest_discover_test testfile.py [output-file.cmake]
#
# This script generates a ctest for each test (functions starting with "test_")
# that are found in the in the given testfile.py.
# The cmake test "add_test" and "set_property" commands are written either
# - to a cmake file (if a second argument output-file.cmake is given). This .cmake file needs to be included
#   in the CMakeLists.txt of the target in order to register the test during configuration
# - to stdout (for test purpose, when there is only one argument)
#
# All generated cmake tests (see below):
# - are run in VENV_DIR virtual environment
# - add an extra PYTHONPATH (using ${PYTHON_EXTRA_PATH})
# - are labelled with a PYTEST_TARGET prefix
#
# WARNING : this script REQUIRES the following variables to be defined before the .cmake output is included:
# - VENV_DIR: the venv directory in which the tests have to be run
# - PYTHON_EXTRA_PATH: a path to a directory that contains a cpython module that is required by the test 
#   (e.g. PYTHON_EXTRA_PATH=build/lib is required to test the CamiTK python module as it contains CamiTK python module shared object)
# - PYTEST_TARGET the name of the target the tests have to be associated with
#
# To test manually:
# - cd build
# - PYTEST_SRC=../sdk/libraries/python/testing; PYTEST_VENV=sdk/libraries/python/testing/.venv; 
# - get the output of pytest_discover_tests:
#   ${PYTEST_VENV}/bin/python ${PYTEST_SRC}/pytest_discover_tests.py ${PYTEST_SRC}/camitk_python_test.py
#
import ast
import sys
import os
from pathlib import Path
from contextlib import contextmanager

@contextmanager
def file_or_stdout(file_name):
    if file_name is None:
        yield sys.stdout
    else:
        with open(file_name, 'w') as out_file:
            yield out_file

def extract_test_functions(filename):
    with open(filename, "r") as f:
        tree = ast.parse(f.read(), filename=filename)
    return [node.name for node in tree.body if isinstance(node, ast.FunctionDef) and node.name.startswith("test_")]

def main():
    if len(sys.argv) < 3 or len(sys.argv) > 5:
        print("Usage: pytest_to_ctest.py path/to/test_file.py [/path/to/output_file.cmake] [PYTHON_PATH] [PATH_env]", file=sys.stderr)
        sys.exit(1)

    test_file = sys.argv[1]
    if not os.path.exists(test_file):
        print(f"Error: File {test_file} does not exist.", file=sys.stderr)
        sys.exit(1)
    
    output_ctest_file = None
    ESCAPED_PATH = None
    PYTHON_PATH = ""
    ENV_PROPERTIES = ""
    if len(sys.argv) >= 3:
        output_ctest_file = sys.argv[2]      
        PYTHON_PATH = sys.argv[3].replace("\\","\\\\").replace(";","$<SEMICOLON>")
        ENV_PROPERTIES = f"PYTHONPATH={PYTHON_PATH}"
        if len(sys.argv) == 5:
            ESCAPED_PATH = sys.argv[4].replace("\\","\\\\").replace(";","$<SEMICOLON>")
            ENV_PROPERTIES += f";PATH={ESCAPED_PATH}"
    test_basename = Path(test_file).stem.replace("_","-")
    
    test_function_names = extract_test_functions(test_file)

    if len(test_function_names) == 0:
        print("Error in listing tests of ", sys.argv[1])
        print("No test function found")
    else:
        with file_or_stdout(output_ctest_file) as f:
            print("Generating to " + output_ctest_file if output_ctest_file is not None else "Generating to stdout")
            for test_name in test_function_names:
                ctest_name = test_basename + "-" + test_name.replace("_","-")[5:]
                print(f"message(STATUS \"Creating test {ctest_name}\")", file=f)
                print(f"add_test(NAME {ctest_name} COMMAND ${{VENV_EXECUTABLE_DIR}}/python -m pytest -vv {test_file} -k {test_name})", file=f)
                # add the extra python path is required to find the CamiTK python module
                # add the working directly (the build directory or install directory = a camitk repository) is required for camitk to load all existing extensions
                print(f"set_tests_properties({ctest_name} PROPERTIES ENVIRONMENT \"{ENV_PROPERTIES}\" WORKING_DIRECTORY ${{CMAKE_BINARY_DIR}} LABELS ${{PYTEST_TARGET}}-{ctest_name})", file=f)

if __name__ == "__main__":
    main()