#!/usr/bin/env python3

import argparse
import os
import re
import subprocess
from pathlib import Path
from typing import Optional, Sequence

parser = argparse.ArgumentParser()
parser.add_argument("--worker-out-dir", default="./test-results", help="output directory")
parser.add_argument("pytest_args", nargs="*", help="args passed to pytest")


def circleci_get_current_bin(args, tests: list[str]) -> list[str]:
    """Get the current bin of tests from CircleCI.

    Uses the CIRCLECI_NODE_INDEX environment variable.
    """
    _ = args
    this_worker_tests = subprocess.run(
        ["circleci", "tests", "split", "--split-by=timings"],
        input="\n".join(tests).encode("utf-8"),
        capture_output=True,
        check=True,
    )
    this_worker_tests_list = this_worker_tests.stdout.decode("utf-8").splitlines()

    return this_worker_tests_list


def collect(pytest_args: list[str]) -> list[str]:
    list_tests_out = subprocess.run(
        ["pytest", "-q", "--collect-only", "--disable-warnings", *pytest_args],
        capture_output=True,
        check=True,
    )
    filter_out = re.compile("(^$|^.* tests collected .*$)")
    tests = [test for test in list_tests_out.stdout.decode("utf-8").splitlines() if not filter_out.match(test)]
    return tests


def circleci_worker(args: argparse.Namespace):
    test_names = collect(args.pytest_args)
    worker_id = int(os.environ["CIRCLE_NODE_INDEX"])
    this_worker_tests = circleci_get_current_bin(args, test_names)

    junitxml = Path(args.worker_out_dir).expanduser() / f"worker-{worker_id}.xml"
    subprocess.run(["pytest", f"--junitxml={junitxml}", "--verbose", "--", *this_worker_tests], check=True)


def main(args: Optional[Sequence[str]] = None):
    parsed_args = parser.parse_args(args)
    circleci_worker(parsed_args)


if __name__ == "__main__":
    main()
