# -*- coding: utf-8 -*- #
# Copyright 2020 Google LLC. All Rights Reserved.
#
# 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.
"""Utilities for flags for `gcloud tasks` commands."""

from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals

from googlecloudsdk.calliope import arg_parsers
from googlecloudsdk.calliope import base
from googlecloudsdk.calliope.concepts import concepts
from googlecloudsdk.calliope.concepts import deps
from googlecloudsdk.command_lib.compute.networks import flags as compute_network_flags
from googlecloudsdk.command_lib.compute.networks.subnets import flags as compute_subnet_flags
from googlecloudsdk.command_lib.kms import resource_args as kms_resource_args
from googlecloudsdk.command_lib.notebooks import completers
from googlecloudsdk.command_lib.util.concepts import concept_parsers
from googlecloudsdk.core import properties


def GetEnvironmentResourceArg(
    api_version, help_text, positional=True, required=True
):
  """Constructs and returns the Environment Resource Argument."""

  def GetEnvironmentResourceSpec():
    """Constructs and returns the Resource specification for Environment."""

    def EnvironmentAttributeConfig():
      return concepts.ResourceParameterAttributeConfig(
          name='environment', help_text=help_text)

    def LocationAttributeConfig():
      return concepts.ResourceParameterAttributeConfig(
          name='{}location'.format('' if positional else 'environment-'),
          help_text=(
              'Google Cloud location of this environment '
              'https://cloud.google.com/compute/docs/regions-zones/#locations.'
          ),
          completer=completers.LocationCompleter,
          fallthroughs=[
              deps.ArgFallthrough('--location'),
              deps.PropertyFallthrough(properties.VALUES.notebooks.location)
          ],
      )

    return concepts.ResourceSpec(
        'notebooks.projects.locations.environments',
        resource_name='environment',
        api_version=api_version,
        environmentsId=EnvironmentAttributeConfig(),
        locationsId=LocationAttributeConfig(),
        projectsId=concepts.DEFAULT_PROJECT_ATTRIBUTE_CONFIG,
        disable_auto_completers=False,
    )

  return concept_parsers.ConceptParser.ForResource(
      '{}environment'.format('' if positional else '--'),
      GetEnvironmentResourceSpec(),
      help_text,
      required=required,
  )


def AddListEnvironmentFlags(parser):
  parser.add_argument(
      '--location',
      completer=completers.LocationCompleter,
      help=(
          'Google Cloud location of this environment: '
          'https://cloud.google.com/compute/docs/regions-zones/#locations.'
      ),
  )


def AddCreateEnvironmentFlags(api_version, parser):
  """Construct groups and arguments specific to the environment create."""
  source_group = parser.add_group(mutex=True, required=True)
  vm_source_group = source_group.add_group()
  vm_mutex_group = vm_source_group.add_group(mutex=True, required=True)
  container_group = source_group.add_group()
  GetEnvironmentResourceArg(
      api_version,
      (
          'User-defined unique name of this environment. The environment name'
          ' must be 1 to 63 characters long and contain only lowercase letters,'
          ' numeric characters, and dashes. The first character must be a'
          ' lowercaseletter and the last character cannot be a dash.'
      ),
  ).AddToParser(parser)
  parser.add_argument(
      '--description', help='A brief description of this environment.'
  )
  parser.add_argument('--display-name', help='Name to display on the UI.')
  parser.add_argument(
      '--post-startup-script',
      help=(
          'Path to a Bash script that automatically runs after a notebook'
          ' instance fully boots up. The path must be a URL or Cloud Storage'
          ' path(gs://`path-to-file/`file-name`).'
      ),
  )
  base.ASYNC_FLAG.AddToParser(parser)
  vm_source_group.add_argument(
      '--vm-image-project',
      help=(
          'The ID of the Google Cloud project that this VM image belongs to.'
          'Format: projects/`{project_id}`.'
      ),
      default='deeplearning-platform-release',
  )
  vm_mutex_group.add_argument(
      '--vm-image-family',
      help=(
          'Use this VM image family to find the image; the newest image in '
          'this family will be used.'
      ),
      default='common-cpu',
  )
  vm_mutex_group.add_argument(
      '--vm-image-name', help='Use this VM image name to find the image.'
  )
  container_group.add_argument(
      '--container-repository',
      help=(
          'The path to the container image repository. For example: '
          'gcr.io/`{project_id}`/`{image_name}`.'
      ),
      required=True,
  )
  container_group.add_argument(
      '--container-tag',
      help=(
          'The tag of the container image. If not specified, this defaults to'
          ' the latest tag.'
      ),
  )


def AddDeleteEnvironmentFlags(api_version, parser):
  GetEnvironmentResourceArg(
      api_version,
      (
          'User-defined unique name of this environment. The environment name'
          ' must be 1 to 63 characters long and contain only lowercase letters,'
          ' numeric characters, and dashes. The first character must be a'
          ' lowercaseletter and the last character cannot be a dash.'
      ),
  ).AddToParser(parser)
  base.ASYNC_FLAG.AddToParser(parser)


def AddDescribeEnvironmentFlags(api_version, parser):
  GetEnvironmentResourceArg(
      api_version,
      (
          'User-defined unique name of this environment. The environment name'
          ' must be 1 to 63 characters long and contain only lowercase letters,'
          ' numeric characters, and dashes. The first character must be a'
          ' lowercaseletter and the last character cannot be a dash.'
      ),
  ).AddToParser(parser)


def GetLocationResourceArg(api_version, help_text):
  """Constructs and returns the Location Resource Argument."""

  def GetLocationResourceSpec():
    """Constructs and returns the Location Resource Argument."""
    def LocationAttributeConfig():
      return concepts.ResourceParameterAttributeConfig(
          name='location',
          help_text=(
              'Google Cloud location of this runtime '
              'https://cloud.google.com/compute/docs/regions-zones/#locations.'
          ),
          fallthroughs=[
              deps.PropertyFallthrough(properties.VALUES.notebooks.location),
          ],
      )

    return concepts.ResourceSpec(
        'notebooks.projects.locations',
        resource_name='location',
        api_version=api_version,
        projectsId=concepts.DEFAULT_PROJECT_ATTRIBUTE_CONFIG,
        locationsId=LocationAttributeConfig())

  return concept_parsers.ConceptParser.ForResource(
      '--location', GetLocationResourceSpec(), help_text, required=True)


def GetRuntimeResourceArg(api_version, help_text):
  """Constructs and returns the Runtime Resource Argument."""

  def GetRuntimeResourceSpec():
    """Constructs and returns the Resource specification for Runtime."""

    def RuntimeAttributeConfig():
      return concepts.ResourceParameterAttributeConfig(
          name='runtime', help_text=help_text)

    def LocationAttributeConfig():
      return concepts.ResourceParameterAttributeConfig(
          name='location',
          help_text=(
              'Google Cloud location of this runtime '
              'https://cloud.google.com/compute/docs/regions-zones/#locations.'
          ),
          fallthroughs=[
              deps.PropertyFallthrough(properties.VALUES.notebooks.location),
          ],
      )

    return concepts.ResourceSpec(
        'notebooks.projects.locations.runtimes',
        resource_name='runtime',
        api_version=api_version,
        projectsId=concepts.DEFAULT_PROJECT_ATTRIBUTE_CONFIG,
        locationsId=LocationAttributeConfig(),
        runtimesId=RuntimeAttributeConfig(),
        disable_auto_completers=False)

  return concept_parsers.ConceptParser.ForResource(
      'runtime', GetRuntimeResourceSpec(), help_text, required=True)


def GetInstanceResourceArg(api_version, help_text):
  """Constructs and returns the Instance Resource Argument."""

  def GetInstanceResourceSpec():
    """Constructs and returns the Resource specification for Instance."""

    def InstanceAttributeConfig():
      return concepts.ResourceParameterAttributeConfig(
          name='instance', help_text=help_text)

    def LocationAttributeConfig():
      return concepts.ResourceParameterAttributeConfig(
          name='location',
          help_text=(
              'Google Cloud location of this environment '
              'https://cloud.google.com/compute/docs/regions-zones/#locations.'
          ),
          fallthroughs=[
              deps.PropertyFallthrough(properties.VALUES.notebooks.location),
          ],
      )

    return concepts.ResourceSpec(
        'notebooks.projects.locations.instances',
        resource_name='instance',
        api_version=api_version,
        instancesId=InstanceAttributeConfig(),
        locationsId=LocationAttributeConfig(),
        projectsId=concepts.DEFAULT_PROJECT_ATTRIBUTE_CONFIG,
        disable_auto_completers=False)

  return concept_parsers.ConceptParser.ForResource(
      'instance', GetInstanceResourceSpec(), help_text, required=True)


def AddNetworkArgument(help_text, parser):
  """Adds Resource arg for network to the parser."""

  def GetNetworkResourceSpec():
    """Constructs and returns the Resource specification for Subnet."""

    def NetworkAttributeConfig():
      return concepts.ResourceParameterAttributeConfig(
          name='network',
          help_text=help_text,
          completer=compute_network_flags.NetworksCompleter)

    return concepts.ResourceSpec(
        'compute.networks',
        resource_name='network',
        network=NetworkAttributeConfig(),
        project=concepts.DEFAULT_PROJECT_ATTRIBUTE_CONFIG,
        disable_auto_completers=False)

  concept_parsers.ConceptParser.ForResource('--network',
                                            GetNetworkResourceSpec(),
                                            help_text).AddToParser(parser)


def AddSubnetArgument(help_text, parser):
  """Adds Resource arg for subnetwork to the parser."""

  def GetSubnetResourceSpec():
    """Constructs and returns the Resource specification for Subnet."""

    def SubnetAttributeConfig():
      return concepts.ResourceParameterAttributeConfig(
          name='subnet',
          help_text=help_text,
          completer=compute_subnet_flags.SubnetworksCompleter)

    def RegionAttributeConfig():
      return concepts.ResourceParameterAttributeConfig(
          name='subnet-region',
          help_text=(
              'Google Cloud region of this subnetwork '
              'https://cloud.google.com/compute/docs/regions-zones/#locations.'
          ),
          completer=completers.RegionCompleter)

    return concepts.ResourceSpec(
        'compute.subnetworks',
        resource_name='subnetwork',
        subnetwork=SubnetAttributeConfig(),
        region=RegionAttributeConfig(),
        project=concepts.DEFAULT_PROJECT_ATTRIBUTE_CONFIG,
        disable_auto_completers=False,
    )

  concept_parsers.ConceptParser.ForResource(
      '--subnet', GetSubnetResourceSpec(), help_text
  ).AddToParser(parser)


def AddRuntimeResource(api_version, parser, add_async_flag=True):
  GetRuntimeResourceArg(
      api_version,
      (
          'User-defined unique name of this runtime. The runtime name must be 1'
          ' to 63 characters long and contain only lowercase letters, numeric'
          ' characters, and dashes. The first character must be a lowercase'
          ' letter and the last character cannot be a dash.'
      ),
  ).AddToParser(parser)
  if add_async_flag:
    base.ASYNC_FLAG.AddToParser(parser)


def AddInstanceResource(api_version, parser, add_async_flag=True):
  GetInstanceResourceArg(
      api_version,
      (
          'User-defined unique name of this instance. The instance name must be'
          ' 1 to 63 characters long and contain only lowercase letters, numeric'
          ' characters, and dashes. The first character must be a lowercase'
          ' letter and the last character cannot be a dash.'
      ),
  ).AddToParser(parser)
  if add_async_flag:
    base.ASYNC_FLAG.AddToParser(parser)


def AddDiagnosticConfigFlags(parser, vm_type):
  """Adds Diagnostic config flags to parser."""
  parser.add_argument(
      '--gcs-bucket',
      help=('The Cloud Storage bucket where the log files generated from the '
            'diagnose command will be stored. storage.buckets.writer '
            'permissions must be given to project\'s service account or '
            'user credential. Format: gs://`{gcs_bucket}` '),
      required=True)
  parser.add_argument(
      '--relative-path',
      help=('Defines the relative storage path in the Cloud Storage bucket '
            'where the diagnostic logs will be written. Default path will be '
            'the root directory of the Cloud Storage bucket'
            'Format of full path: gs://`{gcs_bucket}`/`{relative_path}`/ '),
      required=False)
  parser.add_argument(
      '--enable-repair',
      action='store_true',
      dest='enable-repair',
      default=False,
      help=('Enables flag to repair service for {}'.format(vm_type)),
      required=False)
  parser.add_argument(
      '--enable-packet-capture',
      action='store_true',
      dest='enable-packet-capture',
      default=False,
      help=('Enables flag to capture packets from '
            'the {} for 30 seconds'.format(vm_type)),
      required=False)
  parser.add_argument(
      '--enable-copy-home-files',
      action='store_true',
      dest='enable-copy-home-files',
      default=False,
      help=('Enables flag to copy all `/home/jupyter` folder contents'),
      required=False)
  parser.add_argument(
      '--timeout-minutes',
      help=('Maximum amount of time in minutes before the operation times out'),
      required=False)


def AddCreateInstanceFlags(api_version, parser):
  """Construct groups and arguments specific to the instance creation."""
  accelerator_choices = [
      'NVIDIA_TESLA_A100', 'NVIDIA_TESLA_K80', 'NVIDIA_TESLA_P100',
      'NVIDIA_TESLA_V100', 'NVIDIA_TESLA_P4', 'NVIDIA_TESLA_T4',
      'NVIDIA_TESLA_T4_VWS', 'NVIDIA_TESLA_P100_VWS', 'NVIDIA_TESLA_P4_VWS',
      'TPU_V2', 'TPU_V3', 'NVIDIA_L4', 'NVIDIA_H100_80GB',
      'NVIDIA_H100_MEGA_80GB'
  ]
  disk_choices = ['PD_STANDARD', 'PD_SSD', 'PD_BALANCED', 'PD_EXTREME']
  encryption_choices = ['GMEK', 'CMEK']
  reservation_choices = [
      'TYPE_UNSPECIFIED',
      'NO_RESERVATION',
      'ANY_RESERVATION',
      'SPECIFIC_RESERVATION',
  ]
  AddInstanceResource(api_version, parser)
  environment_group = parser.add_group(mutex=True)
  GetEnvironmentResourceArg(api_version, (
      'User-defined unique name of this environment. The environment name '
      'must be 1 to 63 characters long and contain only lowercase letters, '
      'numeric characters, and dashes. The first character must be a lowercase '
      'letter and the last character cannot be a dash.'),
                            positional=False,
                            required=False).AddToParser(environment_group)
  vm_source_group = environment_group.add_group()
  vm_mutex_group = vm_source_group.add_group(mutex=True)
  container_group = environment_group.add_group()
  vm_source_group.add_argument(
      '--vm-image-project',
      help=('The ID of the Google Cloud project that this VM image belongs to. '
            'Format: projects/`{project_id}`.'),
      default='deeplearning-platform-release')
  vm_mutex_group.add_argument(
      '--vm-image-family',
      help=('Use this VM image family to find the image; the newest image in '
            'this family will be used.'),
      default='common-cpu')
  vm_mutex_group.add_argument(
      '--vm-image-name', help='Use this VM image name to find the image.')
  container_group.add_argument(
      '--container-repository',
      help=('The path to the container image repository. '
            'For example: gcr.io/`{project_id}`/`{image_name}`.'),
      required=True)
  container_group.add_argument(
      '--container-tag',
      help=('The tag of the container image. If not specified, '
            'this defaults to the latest tag.'))
  parser.add_argument(
      '--post-startup-script',
      help=(
          'Path to a Bash script that automatically runs after a notebook '
          'instance fully boots up. The path must be a URL or Cloud Storage '
          'path (gs://`path-to-file`/`file-name`).'))
  parser.add_argument(
      '--service-account',
      help=(
          'The service account on this instance, giving access to other '
          'Google Cloud services. You can use any service account within the '
          'same project, but you must have the service account user permission '
          'to use the instance. If not specified, the [Compute Engine default '
          'service account](/compute/docs/access/service-accounts#default_'
          'service_account) is used.'))
  parser.add_argument(
      '--machine-type',
      help=(
          'The '
          '[Compute Engine machine type](https://cloud.google.com/sdk/gcloud/reference/compute/machine-types) '  # pylint: disable=line-too-long
          'of this instance.'),
      default='n1-standard-1')
  parser.add_argument(
      '--instance-owners',
      help=(
          'The owners of this instance after creation. '
          'Format:`alias@example.com`. Currently supports one owner only.'
          'If not specified, all of the service account users of the VM '
          'instance\'s service account can use the instance.'))
  accelerator_group = parser.add_group(
      help=(
          'The hardware accelerator used on this instance. If you use '
          'accelerators, make sure that your configuration has [enough vCPUs '
          'and memory to support the `machine_type` you have selected]'
          '(/compute/docs/gpus/#gpus-list).'))
  accelerator_group.add_argument(
      '--accelerator-type',
      help='Type of this accelerator.',
      choices=accelerator_choices,
      default=None)
  accelerator_group.add_argument(
      '--accelerator-core-count',
      help='Count of cores of this accelerator.',
      type=int)
  gpu_group = parser.add_group(help='GPU driver configurations.')
  gpu_group.add_argument(
      '--install-gpu-driver',
      action='store_true',
      dest='install_gpu_driver',
      help=(
          'Whether the end user authorizes Google Cloud to install a GPU '
          'driver on this instance. If this field is empty or set to false, '
          'the GPU driver won\'t be installed. Only applicable to instances '
          'with GPUs.'
      ))
  gpu_group.add_argument(
      '--custom-gpu-driver-path',
      help=(
          'Specify a custom Cloud Storage path where the GPU driver is '
          'stored. If not specified, we\'ll automatically choose from official '
          'GPU drivers.'))
  boot_group = parser.add_group(help='Boot disk configurations.')
  boot_group.add_argument(
      '--boot-disk-type',
      choices=disk_choices,
      default=None,
      help=(
          'Type of boot disk attached to this instance, defaults to '
          'standard persistent disk (`PD_STANDARD`).'
      ),
  )
  boot_group.add_argument(
      '--boot-disk-size',
      type=int,
      help=(
          'Size of boot disk in GB attached to this instance, up to '
          'a maximum of 64000 GB (64 TB). The minimum recommended value '
          'is 100 GB. If not specified, this defaults to 100.'
      ),
  )
  data_group = parser.add_group(help='Data disk configurations.')
  data_group.add_argument(
      '--data-disk-type',
      choices=disk_choices,
      default=None,
      help=(
          'Type of data disk attached to this instance, defaults to '
          'standard persistent disk (`PD_STANDARD`).'
      ),
  )
  data_group.add_argument(
      '--data-disk-size',
      type=int,
      help=(
          'Size of data disk in GB attached to this instance, up to '
          'a maximum of 64000 GB (64 TB). The minimum recommended value '
          'is 100 GB. If not specified, this defaults to 100.'
      ),
  )
  data_group.add_argument(
      '--no-remove-data-disk',
      action='store_true',
      dest='no_remove_data_disk',
      help=(
          'If true, the data disk will not be auto deleted when deleting '
          'the instance.'
      ),
  )
  encryption_group = parser.add_group(help='Disk encryption configurations.')
  encryption_group.add_argument(
      '--disk-encryption',
      choices=encryption_choices,
      default=None,
      help='Disk encryption method used on the boot disk, defaults to GMEK.')
  kms_resource_args.AddKmsKeyResourceArg(encryption_group, 'instance')
  network_group = parser.add_group(help='Network configs.')
  network_group.add_argument(
      '--no-public-ip',
      action='store_true',
      dest='no_public_ip',
      help="""\
  If specified, no public IP will be assigned to this instance.""")
  network_group.add_argument(
      '--no-proxy-access',
      action='store_true',
      dest='no_proxy_access',
      help="""\
  If true, the notebook instance will not register with the proxy.""")

  AddNetworkArgument(
      ('The name of the VPC that this instance is in. Format: '
       'projects/`{project_id}`/global/networks/`{network_id}`.'),
      network_group)
  AddSubnetArgument(
      ('The name of the subnet that this instance is in. Format: projects/'
       '`{project_id}`/regions/`{region}`/subnetworks/`{subnetwork_id}`.'),
      network_group)
  parser.add_argument(
      '--labels',
      help=('Labels to apply to this instance. These can be later modified '
            'by the setLabels method.'),
      type=arg_parsers.ArgDict(),
      metavar='KEY=VALUE')
  parser.add_argument(
      '--metadata',
      help=("""\
Custom metadata to apply to this instance.

For example, to specify a Cloud Storage bucket for automatic backup,
you can use the `gcs-data-bucket` metadata tag. Format:
`"--metadata=gcs-data-bucket=``BUCKET''"`.
"""),
      type=arg_parsers.ArgDict(),
      metavar='KEY=VALUE')
  parser.add_argument(
      '--shielded-secure-boot',
      help=('Boot instance with secure boot enabled. '
            'Disabled by default.'),
      dest='shielded_vm_secure_boot',
      action='store_true',
      default=False,
  )
  parser.add_argument(
      '--shielded-vtpm',
      help=('Boot instance with TPM (Trusted Platform Module) enabled.'),
      dest='shielded_vm_vtpm',
      action='store_true',
      default=True,
  )
  parser.add_argument(
      '--shielded-integrity-monitoring',
      help=('Enable monitoring of the boot integrity of the instance.'),
      dest='shielded_vm_integrity_monitoring',
      action='store_true',
      default=True,
  )
  reservation_group = parser.add_group(
      help='Specifies the reservation for the instance.')
  reservation_group.add_argument(
      '--reservation-affinity',
      choices=reservation_choices,
      default='TYPE_UNSPECIFIED',
      help='The type of reservation for the instance.',
  )
  reservation_group.add_argument(
      '--reservation',
      help=('The name of the reservation, required when '
            '`--reservation-affinity=SPECIFIC_RESERVATION`.'))
  parser.add_argument(
      '--tags',
      metavar='TAGS',
      help=('Tags to apply to this instance.'),
      type=arg_parsers.ArgList())


def AddDescribeInstanceFlags(api_version, parser):
  AddInstanceResource(api_version, parser, add_async_flag=False)


def AddDeleteInstanceFlags(api_version, parser):
  AddInstanceResource(api_version, parser)


def AddListInstanceFlags(parser):
  parser.add_argument(
      '--location',
      completer=completers.LocationCompleter,
      help=('Google Cloud location of this environment '
            'https://cloud.google.com/compute/docs/regions-zones/#locations.'))


def AddRegisterInstanceFlags(api_version, parser):
  AddInstanceResource(api_version, parser)


def AddResetInstanceFlags(api_version, parser):
  AddInstanceResource(api_version, parser)


def AddStartInstanceFlags(api_version, parser):
  AddInstanceResource(api_version, parser)


def AddStopInstanceFlags(api_version, parser):
  AddInstanceResource(api_version, parser)


def AddGetHealthInstanceFlags(api_version, parser):
  AddInstanceResource(api_version, parser, add_async_flag=False)


def AddIsUpgradeableInstanceFlags(api_version, parser):
  AddInstanceResource(api_version, parser, add_async_flag=False)


def AddUpgradeInstanceFlags(api_version, parser):
  AddInstanceResource(api_version, parser)


def AddRollbackInstanceFlags(api_version, parser):
  AddInstanceResource(api_version, parser)
  parser.add_argument(
      '--target-snapshot',
      help='The saved snapshot to rollback to',
      required=True)


def AddUpdateInstanceFlags(api_version, parser):
  """Adds accelerator, labels and machine type flags to the parser for update."""
  accelerator_choices = [
      'NVIDIA_TESLA_A100', 'NVIDIA_TESLA_K80', 'NVIDIA_TESLA_P100',
      'NVIDIA_TESLA_V100', 'NVIDIA_TESLA_P4', 'NVIDIA_TESLA_T4',
      'NVIDIA_TESLA_T4_VWS', 'NVIDIA_TESLA_P100_VWS', 'NVIDIA_TESLA_P4_VWS',
      'TPU_V2', 'TPU_V3', 'NVIDIA_L4', 'NVIDIA_H100_80GB',
      'NVIDIA_H100_MEGA_80GB'
  ]
  AddInstanceResource(api_version, parser)
  update_group = parser.add_group(required=True)
  update_group.add_argument(
      '--accelerator-type',
      help='Type of this accelerator.',
      choices=accelerator_choices,
      default=None)
  update_group.add_argument(
      '--accelerator-core-count',
      help='Count of cores of this accelerator.',
      type=int)
  update_group.add_argument(
      '--labels',
      help=('Labels to apply to this instance. '
            'These can be later modified by the setLabels method.'),
      type=arg_parsers.ArgDict(),
      metavar='KEY=VALUE')
  update_group.add_argument(
      '--machine-type',
      help='The [Compute Engine machine type](/compute/docs/machine-types).')


def AddDiagnoseInstanceFlags(api_version, parser):
  """Construct groups and arguments specific to the instance diagnosing."""
  AddInstanceResource(api_version, parser)
  AddDiagnosticConfigFlags(parser, 'instance')


def AddMigrateInstanceFlags(api_version, parser):
  """Construct groups and arguments specific to the instance migration."""
  post_startup_script_option_choices = [
      'POST_STARTUP_SCRIPT_OPTION_UNSPECIFIED',
      'POST_STARTUP_SCRIPT_OPTION_SKIP',
      'POST_STARTUP_SCRIPT_OPTION_RERUN',
      ]
  AddInstanceResource(api_version, parser)
  parser.add_argument(
      '--post-startup-script-option',
      help='// Specifies the behavior of post startup script during migration.',
      choices=post_startup_script_option_choices,
      default='POST_STARTUP_SCRIPT_OPTION_UNSPECIFIED')


def AddDeleteRuntimeFlags(api_version, parser):
  AddRuntimeResource(api_version, parser)


def AddDescribeRuntimeFlags(api_version, parser):
  AddRuntimeResource(api_version, parser, add_async_flag=False)


def AddStartRuntimeFlags(api_version, parser):
  AddRuntimeResource(api_version, parser)


def AddStopRuntimeFlags(api_version, parser):
  AddRuntimeResource(api_version, parser)


def AddSwitchRuntimeFlags(api_version, parser):
  """Adds accelerator and machine type flags to the parser for switch."""
  accelerator_choices = [
      'NVIDIA_TESLA_A100', 'NVIDIA_TESLA_K80', 'NVIDIA_TESLA_P100',
      'NVIDIA_TESLA_V100', 'NVIDIA_TESLA_P4', 'NVIDIA_TESLA_T4',
      'NVIDIA_TESLA_T4_VWS', 'NVIDIA_TESLA_P100_VWS', 'NVIDIA_TESLA_P4_VWS',
      'TPU_V2', 'TPU_V3'
  ]
  AddRuntimeResource(api_version, parser)
  parser.add_argument('--machine-type', help=('machine type'))
  accelerator_config_group = parser.add_group()
  accelerator_config_group.add_argument(
      '--accelerator-type',
      help='Type of this accelerator.',
      choices=accelerator_choices,
      default=None)
  accelerator_config_group.add_argument(
      '--accelerator-core-count',
      help='Count of cores of this accelerator.',
      type=int)


def AddResetRuntimeFlags(api_version, parser):
  AddRuntimeResource(api_version, parser)


def AddCreateRuntimeFlags(api_version, parser):
  """Construct groups and arguments specific to the runtime creation."""

  AddRuntimeResource(api_version, parser)
  runtime_type_group = parser.add_group(mutex=True, required=True)
  runtime_type_group.add_argument('--runtime-type', help='runtime type')
  machine_type_group = runtime_type_group.add_group()
  machine_type_group.add_argument(
      '--machine-type', help='machine type', required=True
  )
  local_disk_group = machine_type_group.add_group()
  local_disk_group.add_argument('--interface', help='runtime interface')
  local_disk_group.add_argument('--source', help='runtime source')
  local_disk_group.add_argument('--mode', help='runtime mode')
  local_disk_group.add_argument('--type', help='runtime type')

  access_config_group = parser.add_group(required=True)
  access_config_group.add_argument('--runtime-access-type', help='access type')
  access_config_group.add_argument('--runtime-owner', help='runtime owner')

  software_config_group = parser.add_group()
  software_config_group.add_argument(
      '--idle-shutdown-timeout', help='idle shutdown timeout'
  )
  software_config_group.add_argument(
      '--install-gpu-driver', help='install gpu driver'
  )
  software_config_group.add_argument(
      '--custom-gpu-driver-path', help='custom gpu driver path'
  )
  software_config_group.add_argument(
      '--post-startup-script', help='post startup script'
  )
  software_config_group.add_argument(
      '--post-startup-script-behavior', help='post startup script behavior'
  )


def AddListRuntimeFlags(api_version, parser):
  GetLocationResourceArg(
      api_version, 'Location of this runtime. For example, us-central1-a'
  ).AddToParser(parser)


def AddDiagnoseRuntimeFlags(api_version, parser):
  """Construct groups and arguments specific to the runtime diagnosing."""
  AddRuntimeResource(api_version, parser)
  AddDiagnosticConfigFlags(parser, 'runtime')


def AddMigrateRuntimeFlags(api_version, parser):
  """Construct groups and arguments specific to the runtime migration."""
  post_startup_script_option_choices = [
      'POST_STARTUP_SCRIPT_OPTION_UNSPECIFIED',
      'POST_STARTUP_SCRIPT_OPTION_SKIP',
      'POST_STARTUP_SCRIPT_OPTION_RERUN',
      ]
  AddRuntimeResource(api_version, parser)
  network_group = parser.add_group(help='Network configs.')
  AddNetworkArgument(
      ('The name of the VPC that this instance is in. Format: '
       'projects/`{project_id}`/global/networks/`{network_id}`.'),
      network_group)
  AddSubnetArgument(
      ('The name of the subnet that this instance is in. Format: projects/'
       '`{project_id}`/regions/`{region}`/subnetworks/`{subnetwork_id}`.'),
      network_group)
  parser.add_argument(
      '--service-account',
      help=(
          'The service account to be included in the Compute Engine instance '
          'of the new Workbench Instance when the Runtime uses single user '
          'only mode for permission. If not specified, the [Compute Engine '
          'default service account](https://cloud.google.com/compute/docs/'
          'access/service-accounts#default_service_account) is used. When the '
          'Runtime uses service account mode for permission, it will reuse the '
          'same service account, and this field must be empty.'
          ))
  parser.add_argument(
      '--post-startup-script-option',
      help='Specifies the behavior of post startup script during migration.',
      choices=post_startup_script_option_choices,
      default='POST_STARTUP_SCRIPT_OPTION_UNSPECIFIED')