start

#!/usr/bin/env python3

# region copyright
# Copyright 2023 NVIDIA Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# endregion


import os
import sys
import time
from pathlib import Path

import click

from src.python.alicloud import (
    alicloud_configure_cli,
    alicloud_get_instance_status,
    alicloud_start_instance,
)
from src.python.aws import (
    aws_configure_cli,
    aws_get_instance_status,
    aws_start_instance,
)
from src.python.azure import azure_login, azure_start_instance
from src.python.config import c as config
from src.python.utils import (
    colorize_error,
    colorize_info,
    colorize_result,
    deployments,
    format_app_name,
    format_cloud_name,
    read_tf_output,
    shell_command,
)

APP_NAMES = [
    "isaac",
]


@click.command()
@click.option(
    "--debug/--no-debug",
    default=False,
    show_default=True,
)
@click.argument(
    "deployment_name",
    required=True,
    type=click.Choice(deployments()),
)
@click.option(
    "--app",
    default="all",
    type=click.Choice(["all"] + APP_NAMES),
    help="Application.",
)
@click.argument(
    "app",
    required=False,
)
def main(
    app: str,
    debug: bool,
    deployment_name: str,
):
    # cleck what cloud is deployment for
    cloud = read_tf_output(deployment_name, "cloud", verbose=debug)

    # go through apps

    apps = APP_NAMES if app == "all" else [app]

    for app_name in apps:
        vm_id = read_tf_output(deployment_name, f"{app_name}_vm_id", verbose=debug)
        vm_ip = read_tf_output(deployment_name, f"{app_name}_ip", verbose=debug)

        if vm_id in ["NA", ""]:
            # show "not found" message if:
            # - app is explicitly specified
            # - debug is enabled
            if app != "all" or debug:
                click.echo(colorize_info(f"* No {format_app_name(app_name)} VM found"))
            continue

        click.echo(
            colorize_info(
                f"* Starting {format_app_name(app_name)} VM @ {format_cloud_name(cloud)} ({vm_ip})..."
            ),
            nl=False,
        )

        # start VM

        if "azure" == cloud:
            click.echo()  # extra new line
            # log into azure if needed
            azure_login(verbose=debug)
            # start
            azure_start_instance(vm_id, verbose=debug)

        elif "alicloud" == cloud:
            # configure aliyun cli
            alicloud_configure_cli(deployment_name, verbose=debug)
            # start
            alicloud_start_instance(vm_id, verbose=debug)
            # wait for VM to stop
            while True:
                if "Running" == alicloud_get_instance_status(vm_id, verbose=debug):
                    click.echo()
                    break
                click.echo(colorize_info("."), nl=False)
                time.sleep(2)

        elif "aws" == cloud:
            # configure aws cli
            aws_configure_cli(deployment_name, verbose=debug)
            # start
            aws_start_instance(vm_id, verbose=debug)
            # wait for VM to start
            while True:
                status = aws_get_instance_status(vm_id, verbose=debug)
                if "running" == status:
                    click.echo()
                    break
                click.echo(colorize_info("."), nl=False)
                time.sleep(2)
        else:
            click.echo(
                colorize_error(
                    f"* Stop/start is not yet supported for {format_cloud_name(cloud)}."
                )
            )
            exit(1)

        click.echo(colorize_result(f"* {format_app_name(app_name)} VM started."))

        # run "on_stop_start" tags in ansible
        click.echo(colorize_info(f"* Configuring {format_app_name(app_name)} VM..."))
        shell_command(
            f"ansible-playbook -i {config['state_dir']}/{deployment_name}/.inventory"
            + f" {app_name}.yml"
            + " -t on_stop_start"
            + f" {'-vv' if debug else ''}",
            cwd=config["ansible_dir"],
            verbose=debug,
        )

        text = Path(f"{config['state_dir']}/{deployment_name}/info.txt").read_text()
        click.echo(colorize_result(text))


if __name__ == "__main__":
    if os.path.exists("/.dockerenv"):
        # we're in docker, run command
        main()
    else:
        # we're outside, start docker container first
        shell_command(f"./run '{' '.join(sys.argv)}'", verbose=True)