blueshift

#!/usr/bin/env python3
import argparse
import os
import sys
import atexit
import subprocess

PROJECT_NAME = os.environ.get("PROJECT_NAME")
PROJECT_PREFIX = os.environ.get("PROJECT_PREFIX")
NETWORK_NAME = os.environ.get("NETWORK_NAME")
DEFAULT_NUM_CLIENTS = os.environ.get("NUM_CLIENTS")

def supressPythonErrors():
    pass

def clean_shutdown(base_cmd):
    subprocess.run(base_cmd + ["down"])

def run(args):
    sys.stderr = supressPythonErrors()
    if args.local:
        bin_dir = os.environ.get("RUNTIME_OUTPUT_DIRECTORY")
        executable = os.path.join("." ,"build", bin_dir, args.binary)
        subprocess.run([executable, *args.binary_args], check=True)
    else:
        subprocess.run(["docker", "compose", "run", "--rm", "builder",  "run", "-l", args.binary, *args.binary_args], check=True)

def deploy(args):
    base_cmd = ["docker", "compose"]
    atexit.register(clean_shutdown, base_cmd)
    
    command = [*base_cmd, "up"]
    if args.num_clients:
        os.environ["NUM_CLIENTS"] = args.num_clients
    if args.build:
        command.append("--build")

    subprocess.run(command)

def build(args):
    sys.stderr = supressPythonErrors()
    build_args = []
    if args.image_build:
        subprocess.run(["docker", "compose", "build"])
    if args.clean:
        build_args.append("--clean")
        subprocess.run(["rm", "-rf", "build"])
    
    if not args.no_local_generation:
        subprocess.run(["cmake", "-S", ".", "-B", "build"], check=True)

    if args.local:
        subprocess.run(["make", *args.make], cwd="./build", check=True)
    else:
        subprocess.run(["docker", "compose", "run", "--rm", "builder",  "build", "-l", *build_args, *args.make], check=True)

def test(args):
    sys.stderr = supressPythonErrors()
    os.environ["GTEST_COLOR"] = "1"
    ctest_args = [args.verbose, "-R", args.tests_regex]
    ctest_args = list(filter(None, ctest_args))
    if args.local:
        fail_output = ["--output-on-failure"] if not args.no_output_on_failure else []
        subprocess.run(["ctest", *fail_output, *ctest_args], cwd="./build", check=True)
    else:
        subprocess.run(["docker", "compose", "run", "--rm", "builder", "test", "-l", *ctest_args], check=True)

def check_successful_reset_or_force():
    network_id = subprocess.check_output(f"docker network ls | grep {NETWORK_NAME} | awk '{{ print $1 }}'",
                                         text=True, shell=True)

    container_ids = subprocess.check_output(f"docker container ls | grep {PROJECT_PREFIX} | awk '{{ print $1 }}'",
                                            text=True, shell=True).strip()

    image_ids = subprocess.check_output(f"docker image ls | grep {PROJECT_PREFIX} | awk '{{ print $3 }}'",
                                     text=True, shell=True).strip()
    
    volumes = subprocess.check_output(f"docker volume ls | grep {PROJECT_NAME} | awk '{{ print $2 }}'",
                                     text=True, shell=True).strip()
        
    if any((network_id, container_ids, image_ids, volumes)):
        subprocess.run(["docker", "network", "rm", network_id])
        subprocess.run(["docker", "container", "stop"] + container_ids.split('\n'))
        subprocess.run(["docker", "container", "rm", "-f"] + container_ids.split('\n'))
        subprocess.run(["docker", "image", "prune", "-f"])
        subprocess.run(["docker", "image", "rm", "-f"] + image_ids.split('\n'))
        subprocess.run(["docker", "volume", "rm", "-f"] + volumes.split('\n'))

def reset(args):
    subprocess.run(["docker", "compose", "down", "--rmi", "all", "-v", "--remove-orphans"])
    subprocess.run(["rm", "-rf", "build"])
    
    check_successful_reset_or_force()

    if args.rm_build_cache:
        subprocess.run(["docker", "buildx", "prune", "-f"])

parser = argparse.ArgumentParser(description=f"{PROJECT_NAME} development environment build script", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
subparsers = parser.add_subparsers(title="commands")

run_parser = subparsers.add_parser("run", help=f"run selected {PROJECT_NAME} binaries")
run_parser.add_argument("binary",
                        help="binary to execute")
run_parser.add_argument("binary_args",
                        help="program arguments",
                        nargs="*",
                        default=None)
run_parser.add_argument("-l", "--local",
                        help="run selected binary on local host instead of in containerized environment",
                        action="store_true")
run_parser.set_defaults(fn=run)

deploy_parser = subparsers.add_parser("deploy", help=f"run {PROJECT_NAME} deployment environment (send SIGINT [ctrl+c] to stop gracefully)")
deploy_parser.add_argument("-n", "--num-clients",
                           help="number of clients to create (only used in deployment simulation)",
                           default=DEFAULT_NUM_CLIENTS)
deploy_parser.add_argument("-b", "--build",
                           help="rebuild container images before running",
                           action="store_true")
deploy_parser.set_defaults(fn=deploy)

build_parser = subparsers.add_parser("build", help=f"build {PROJECT_NAME} binaries or deployment images")
build_parser.add_argument("make",
                          help="makefile arguments",
                          nargs="*",
                          default=None)
build_parser.add_argument("-i", "--image-build",
                          help="build container images",
                          action="store_true")
build_parser.add_argument("-l", "--local",
                          help="build for local host only; container builds are not updated",
                          action="store_true")
build_parser.add_argument("-c", "--clean",
                          help="clean build directory before building",
                          action="store_true")
build_parser.add_argument("--no-local-generation",
                          help="disable generation of build artifacts in local environment",
                          action="store_true")
build_parser.set_defaults(fn=build)

test_parser = subparsers.add_parser("test", help=f"test {PROJECT_NAME} binary and library builds")
test_parser.add_argument("-v", "--verbose",
                         help="show verbose output",
                         action="store_const",
                         const="--verbose")
test_parser.add_argument("--no-output-on-failure",
                         help="show condensed output regardless of failed tests",
                         action="store_true")
test_parser.add_argument("-l", "--local",
                         help="test on local host only; container builds are not tested",
                         action="store_true")
test_parser.add_argument("-R", "--tests-regex",
                         help="run subset of tests matching regular expression",
                         default=".*")
test_parser.set_defaults(fn=test)

reset_parser = subparsers.add_parser("reset", help="resets images and volumes")
reset_parser.add_argument("--rm-build-cache",
                          help="clears dangling docker build cache",
                          action="store_true")
reset_parser.set_defaults(fn=reset)

args = parser.parse_args(args=None if len(sys.argv) > 1 else ["--help"])
args.fn(args)