stop

#!/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

import click

from src.python.alicloud import (
    alicloud_configure_cli,
    alicloud_get_instance_status,
    alicloud_stop_instance,
)
from src.python.aws import aws_configure_cli, aws_get_instance_status, aws_stop_instance
from src.python.azure import azure_login, azure_stop_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

        # stop VM

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

        extra_info = ""

        if "azure" == cloud:
            click.echo()  # extra newline
            extra_info = " On Azure stopped VMs still count towards the quota."
            # login to azure if needed
            azure_login(verbose=debug)
            # stop
            azure_stop_instance(vm_id, verbose=debug)

        elif "alicloud" == cloud:
            # configure aliyun cli
            alicloud_configure_cli(deployment_name, verbose=debug)
            # stop
            alicloud_stop_instance(vm_id, verbose=debug)
            # wait for vm to stop
            while True:
                if "Stopped" == 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)
            # stop
            aws_stop_instance(vm_id, verbose=debug)
            # wait for vm to stop
            while True:
                status = aws_get_instance_status(vm_id, verbose=debug)
                if "stopped" == 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 stopped."
                + " Please note: Storage and network resources may still incur costs."
                + extra_info
            )
        )


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)