import sys
sys.path.append('.')
sys.path.append('../')
import unittest
import argparse
import os
import glob

# Define the test file names
test_files = [
    "testAdvancedFeatures.py",
    "testActivationFunctions.py",
    "testActions.py",
    "testEdgeCases.py",
    "testHistory.py",
    "testLayer.py",
    "testLRscheduler.py",
    "testModel.py",
    "testNewResidualConenctions.py",
    "testOptimizers.py",
    "testPerformance.py",
    "testQIdentity.py",
    "testSimulationScheduler.py",
    "testSimulationScore.py",
    "testStorage.py",
    "testStopper.py",
    "testTrain.py",
    "testTrainer.py",
    "testTrainingScenarios.py",
    "unittests_autogenerated.py",
    "unittests_basic.py",
    "test_neuron_reduction.py",
    "test_neuron_addition.py"
]


def cleanup_test_files():
    """Clean up any test files created during test execution"""
    # List of file patterns to clean up
    cleanup_patterns = [
        "testimage*.html",
        "*.png",
        "result/*",
        "test_*.json",
        "test_*.pkl",
        "resultGNN_model*.html",
        "resultGNN_model*.json",
        "resultGNN_model_hist*"
    ]
    
    for pattern in cleanup_patterns:
        for file in glob.glob(pattern):
            try:
                if os.path.isfile(file):
                    os.remove(file)
                elif os.path.isdir(file):
                    for root, dirs, files in os.walk(file, topdown=False):
                        for name in files:
                            os.remove(os.path.join(root, name))
                        for name in dirs:
                            os.rmdir(os.path.join(root, name))
                    os.rmdir(file)
            except Exception as e:
                print(f"Warning: Could not remove {file}: {str(e)}")

# CPU-only mode - no GPU support
parser = argparse.ArgumentParser(description="Run test suite (CPU only).")
parser.add_argument('--discover', action='store_true', help="Automatically discover all test files in the tests directory")
args = parser.parse_args()

# Create a test loader
loader = unittest.TestLoader()

# Create an empty test suite
suite = unittest.TestSuite()

if args.discover:
    # Automatically discover all test files in the tests directory
    print("Discovering all test files in the tests directory...")
    discovered_tests = loader.discover('./tests/', pattern='test*.py')
    suite.addTests(discovered_tests)
    print(f"Discovered {suite.countTestCases()} test cases")
else:
    # Add tests from each test file to the test suite
    for test_file in test_files:
        # Discover tests in the test file
        discovered_tests = loader.discover('./tests/', pattern=test_file)
        suite.addTests(discovered_tests)

# Create a test runner that will output the results to the console
runner = unittest.TextTestRunner(verbosity=2)

# Run the test suite and exit with appropriate exit code
if __name__ == '__main__':
    try:
        result = runner.run(suite)
        if result.wasSuccessful():
            sys.exit(0)  # Exit with code 0 if all tests passed successfully
        else:
            sys.exit(1)  # Exit with code 1 if there were failures or errors in tests
    finally:
        cleanup_test_files()  # Clean up test files regardless of test outcome
