import signal
import sys
import toml

from docker import rebuild_images_if_needed, Compose
from pathlib import Path
from utils import *

CARGO_INTEGRATION_TOML = Path('../__Cargo_integration.toml')
CLUSTER = Compose()


def fail(msg):
    err(msg)
    sys.exit(1)


def usage_exit():
    usage = f"""Usage:
    python {__file__.split('/')[-1]} [--rebuild] [--preserve] [--no-run] [--down]"""

    fail(usage)


def create_integration_workspace():
    if CARGO_INTEGRATION_TOML.exists():
        CARGO_INTEGRATION_TOML.unlink()
    workspace = toml.load('../Cargo.toml')
    workspace['workspace']['members'] = ['integration']
    with open(CARGO_INTEGRATION_TOML, 'w') as fo:
        toml.dump(workspace, fo)


def run_tests():
    allowed_args = set(["--rebuild", "--preserve", "--no-run", "--release", "--down"])
    args = sys.argv[1:]
    if not set(args).issubset(allowed_args):
        usage_exit()
    force_rebuild = '--rebuild' in args
    preserve_containers = '--preserve' in args
    only_setup_cluster = '--no-run' in args
    down_cluster = "--down" in args

    def _cleanup():
        if not preserve_containers and not only_setup_cluster:
            CLUSTER.down()
        CARGO_INTEGRATION_TOML.unlink(missing_ok=True)

    def abort_handler(s, _):
        warn(f'Received signal: {s}, gracefully stopping...')
        _cleanup()

    if down_cluster:
        _cleanup()
        return
        
    for s in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP):
        signal.signal(s, abort_handler)
    rebuild_images_if_needed(force_rebuild)
    create_integration_workspace()
    try:
        CLUSTER.up()
        CLUSTER.is_alive()
        if not only_setup_cluster:
            CLUSTER.run('cargo test --test integration')
    except Exception as e:
        CLUSTER.print_containers_logs()
        fail(e)
    finally:
        _cleanup()


if __name__ == '__main__':
    run_tests()