start

#!/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
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_validate_credentials,
    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_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()),
)
@click.option(
    "--quick/--no-quick",
    "-q/-Q",
    "quick",
    default=False,
    show_default=True,
    help="Quick start: skip full Ansible, only run autorun script.",
)
@click.option(
    "--ansible/--no-ansible",
    "-a/-A",
    "ansible",
    default=True,
    show_default=True,
    hidden=True,
    help="Run Ansible? (--no-ansible is equivalent to --quick.)",
)
def main(
    debug: bool,
    deployment_name: str,
    quick: bool,
    ansible: bool,
):
    # --no-ansible is a synonym for --quick
    if not ansible:
        quick = True

    # 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)

    # go through instance roles

    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)

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

        # start VM

        if "azure" == cloud:
            click.echo()  # extra new line
            # 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:
            # 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_instance_role(instance_role)} VM started.")
        )

        # run ansible
        ansible_tags = "__autorun" if quick else "on_stop_start"
        if quick:
            click.echo(
                colorize_info(
                    f"* Quick start: running autorun only for {format_instance_role(instance_role)} VM..."
                )
            )
        else:
            click.echo(
                colorize_info(
                    f"* Configuring {format_instance_role(instance_role)} VM..."
                )
            )
        shell_command(
            f"ansible-playbook -i {config['state_dir']}/{deployment_name}/.inventory"
            + " isaac-workstation.yaml"
            + f" -t {ansible_tags}"
            + 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)