repair

#!/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 re
import sys
from pathlib import Path

import click

from src.python.aws import aws_validate_credentials
from src.python.azure import azure_login
from src.python.config import c as config
from src.python.utils import (
    colorize_error,
    colorize_info,
    deployments,
    format_cloud_name,
    gcp_login,
    get_my_public_ip,
    read_meta,
    read_tf_output,
    shell_command,
    subnet_from_ip,
)


@click.command()
@click.option(
    "--debug/--no-debug",
    default=False,
    show_default=True,
)
@click.argument(
    "deployment_name",
    required=True,
    type=click.Choice(deployments()),
)
@click.option(
    "--terraform/--no-terraform",
    "-t/-T",
    default=True,
    show_default=True,
    help="Run Teraform?",
)
@click.option(
    "--ansible/--no-ansible",
    "-a/-A",
    default=True,
    show_default=True,
    help="Run ansible?",
)
@click.option(
    "--ingress-cidrs",
    default=None,
    show_default=True,
    help='Update ingress CIDR blocks (comma-separated). Use "myip" for current IP, "myip/16", "myip/24" for subnets.',
)
def main(
    debug: bool,
    deployment_name: str,
    terraform: bool,
    ansible: bool,
    ingress_cidrs: str,
):
    meta = read_meta(deployment_name, debug)
    cloud = read_tf_output(deployment_name, "cloud", verbose=debug)

    click.echo(
        colorize_info(
            f"* Repairing deployment '{deployment_name}' @ {format_cloud_name(cloud)}..."
        )
    )

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

    # update ingress CIDRs in tfvars if requested
    if ingress_cidrs is not None:
        resolved_cidrs = _resolve_ingress_cidrs(ingress_cidrs, debug)
        tfvars_file = f"{config['state_dir']}/{deployment_name}/.tfvars"
        _update_tfvars_ingress_cidrs(tfvars_file, resolved_cidrs, debug)
        if not terraform:
            click.echo(
                colorize_info(
                    "* --ingress-cidrs requires Terraform, enabling --terraform."
                )
            )
            terraform = True

    # pin AMI for AWS to prevent instance replacement from newer base images
    if cloud == "aws" and terraform:
        current_ami = read_tf_output(
            deployment_name, "isaac_workstation_ami_id", verbose=debug
        )
        ami_name = read_tf_output(
            deployment_name, "isaac_workstation_ami_name", verbose=debug
        )
        if current_ami and current_ami != "NA":
            click.echo(
                colorize_info(f"* Pinning base image: {current_ami} ({ami_name})")
            )
            tfvars_file = f"{config['state_dir']}/{deployment_name}/.tfvars"
            _update_tfvars_value(tfvars_file, "ami_id", current_ami, debug)

    # run terraform
    if terraform:
        click.echo(colorize_info("* Re-running Terraform..."))
        tfstate_file = Path(
            f"{config['state_dir']}/{deployment_name}/.tfstate"
        ).absolute()
        shell_command(
            f"terraform init -upgrade -no-color -input=false -reconfigure -backend-config=\"path={tfstate_file}\" {' > /dev/null' if not debug else ''}",
            cwd=f"{config['terraform_dir']}/{cloud}",
            verbose=debug,
        )
        shell_command(
            "terraform apply -auto-approve "
            + f"-var-file={config['state_dir']}/{deployment_name}/.tfvars",
            cwd=f"{config['terraform_dir']}/{cloud}",
            verbose=debug,
        )

    # run ansible
    if ansible:
        click.echo(colorize_info("* Re-running Ansible..."))
        shell_command(
            f"ansible-playbook -i {config['state_dir']}/{deployment_name}/.inventory "
            + f"isaac-workstation.yaml {'-vv' if debug else ''}",
            cwd=f"{config['ansible_dir']}",
            verbose=debug,
        )


def _resolve_ingress_cidrs(ingress_cidrs_str: str, debug: bool) -> list:
    """Resolve special CIDR values (myip, myip/16, etc.) to actual CIDRs."""

    cidrs = [x.strip().lower() for x in ingress_cidrs_str.split(",") if x.strip()]

    # validate
    for cidr in cidrs:
        if cidr not in (
            "auto",
            "myip",
            "myip/8",
            "myip/16",
            "myip/24",
            "mynet",
        ) and not re.match(r"^\d{1,3}(\.\d{1,3}){3}/\d{1,2}$", cidr):
            raise click.BadParameter(colorize_error(f"Invalid CIDR block: {cidr}"))

    resolved = []
    for cidr in cidrs:
        if cidr in ("auto", "myip"):
            cidr = get_my_public_ip(verbose=debug) + "/32"
        elif cidr in ("mynet", "myip/16"):
            cidr = subnet_from_ip(get_my_public_ip(verbose=debug), "16")
        elif cidr == "myip/24":
            cidr = subnet_from_ip(get_my_public_ip(verbose=debug), "24")
        elif cidr == "myip/8":
            cidr = subnet_from_ip(get_my_public_ip(verbose=debug), "8")
        resolved.append(cidr)

    resolved = list(set(resolved))

    if debug:
        click.echo(colorize_info(f"* Resolved ingress CIDRs: {resolved}"))

    return resolved


def _update_tfvars_ingress_cidrs(tfvars_file: str, cidrs: list, debug: bool):
    """Update ingress_cidrs in an existing .tfvars file."""

    if not os.path.exists(tfvars_file):
        raise click.ClickException(f"tfvars file not found: {tfvars_file}")

    new_value = str(cidrs).replace("'", '"')
    _update_tfvars_value(tfvars_file, "ingress_cidrs", new_value, debug, raw=True)

    click.echo(colorize_info(f"* Updated ingress CIDRs to: {cidrs}"))


def _update_tfvars_value(
    tfvars_file: str, key: str, value: str, debug: bool, raw: bool = False
):
    """Update or add a single value in a .tfvars file.
    If raw=True, value is written as-is (for lists). Otherwise quoted as a string."""

    if not os.path.exists(tfvars_file):
        raise click.ClickException(f"tfvars file not found: {tfvars_file}")

    lines = Path(tfvars_file).read_text().splitlines()
    new_lines = []
    found = False
    formatted = value if raw else f'"{value}"'

    for line in lines:
        if line.strip().startswith(key + " ") or line.strip().startswith(key + "="):
            new_lines.append(f"{key} = {formatted}")
            found = True
        else:
            new_lines.append(line)

    if not found:
        new_lines.append(f"{key} = {formatted}")

    Path(tfvars_file).write_text("\n".join(new_lines) + "\n")

    if debug:
        click.echo(colorize_info(f"* Updated {key} = {formatted} in {tfvars_file}"))


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)