#!/usr/bin/env python3

# This script is a wrapper around pytest which performs the following tasks:
# 1) If the IOC in ioc/startioc.py is not running, start it
# 2) Run tests in parallel by starting a separate pytest process
# for each separate motor given in FOLDER. This allows to run tests on different
# motors in parallel, cutting down execution time and testing parallelism.
# The script accepts any arguments and forwards them to the pytest call.
# If a folder is given as an argument, all tests which are not within this
# folder are ignored:
# - `pytestpar` will run all tests parallelized per motor
# - `pytestpar tests/sinqMotor/turboPmac` will run all Turbo PMAC tests parallelized per motor
# - `pytestpar tests/sinqMotor/turboPmac/ax1` will only run the tests for motor `ax1`.
# These tests are run sequentially.

import sys
import subprocess
import os
from pathlib import Path
import time

from tests.conftest import check_ioc_running
from src.classes import IocNotRunning

# Time we excpect the IOC needs to start up
TIMEOUT_IOC_STARTUP = 10

# Separate path-like arguments from other pytest args
path_args = []
extra_args = []

for arg in sys.argv[1:]:
    if '/' in arg or os.path.exists(arg):
        path_args.append(arg)
    else:
        extra_args.append(arg)

# Find all axes paths (pa) within all controller paths (pc)
folders = []
for pc in Path("tests").glob("*"):
    if pc.is_dir() and "__pycache__" not in pc.parts:
        for pa in Path(pc).glob("*"):
            if pa.is_dir() and "__pycache__" not in pa.parts:
                # Ignore all axis paths which start with an underscore
                if not pa.parts[-1].startswith("_"):
                    folders.append(pa)

# Filter folders to run based on path_args
enabled_folders = []
if not path_args:
    enabled_folders = folders
else:
    for folder in folders:
        if any(Path(folder).resolve().as_posix().startswith(Path(arg).resolve().as_posix()) for arg in path_args):
            enabled_folders.append(folder)

if not enabled_folders:
    print(f"No tests were found within the specified path {path_args}")

# Check if the IOC is running and try to start it, if it isn't
try:
    check_ioc_running()
except IocNotRunning:
    print('Starting IOC ...')
    # Start the IOC as a separate process, wait a bit and try again
    path = Path(__file__).resolve().parent / 'ioc' / 'startioc'
    p = subprocess.Popen([path],
                         stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE)
    now = time.time()
    while now + TIMEOUT_IOC_STARTUP > time.time():
        time.sleep(0.5)
        try:
            check_ioc_running()
        except IocNotRunning:
            pass
    check_ioc_running()
    print(
        f'Started IOC successfully (process ID {p.pid}). It will automatically terminate after all tests are done.')

# Run each enabled folder's relevant tests in parallel
processes = []
for folder in enabled_folders:
    command = ["pytest"] + [str(folder)] + extra_args
    proc = subprocess.Popen(command)
    processes.append(proc)

# Wait for all processes and collect exit codes
exit_code = 0
for proc in processes:
    code = proc.wait()
    if code != 0:
        exit_code = code
sys.exit(exit_code)