stop

#!/usr/bin/env python3

# region copyright
# Copyright 2023-2026 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_validate_credentials,
    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_instance_role,
    format_cloud_name,
    read_meta,
    read_tf_output,
    shell_command,
)


@click.command()
@click.option(
    "--debug/--no-debug",
    default=False,
    show_default=True,
)
@click.argument(
    "deployment_name",
    required=True,
    type=click.Choice(deployments()),
)
def main(
    debug: bool,
    deployment_name: str,
):
    # check what cloud is deployment for
    cloud = read_tf_output(deployment_name, "cloud", verbose=debug)

    # validate credentials before any cloud operations
    if "aws" == cloud:
        meta = read_meta(deployment_name, debug)
        region = meta.get("input_params", {}).get("region", None)
        aws_validate_credentials(deployment_name, region=region, verbose=debug)
    elif "azure" == cloud:
        azure_login(verbose=debug)

    for instance_role in ["isaac_workstation"]:
        vm_id = read_tf_output(deployment_name, f"{instance_role}_vm_id", verbose=debug)
        vm_ip = read_tf_output(deployment_name, f"{instance_role}_ip", verbose=debug)

        # stop VM

        click.echo(
            colorize_info(
                f"* Stopping {format_instance_role(instance_role)} 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."
            # 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:
            # 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_instance_role(instance_role)} 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)