# -*- 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.
"""Flags and helpers for the migration jobs related commands."""


import enum

from googlecloudsdk.calliope import arg_parsers


@enum.unique
class ApiType(enum.Enum):
  """This API type is used to differentiate between the classification types of Create requests and Update requests."""
  CREATE = 'create'
  UPDATE = 'update'


def AddNoAsyncFlag(parser):
  """Adds a --no-async flag to the given parser."""
  help_text = (
      'Waits for the operation in progress to complete before returning.'
  )
  parser.add_argument('--no-async', action='store_true', help=help_text)


def AddDisplayNameFlag(parser):
  """Adds a --display-name flag to the given parser."""
  help_text = """
    A user-friendly name for the migration job. The display name can include
    letters, numbers, spaces, and hyphens, and must start with a letter.
    """
  parser.add_argument('--display-name', help=help_text)


def AddTypeFlag(parser, required=False):
  """Adds --type flag to the given parser."""
  help_text = 'Type of the migration job.'
  choices = ['ONE_TIME', 'CONTINUOUS']
  parser.add_argument(
      '--type', help=help_text, choices=choices, required=required
  )


def AddDumpGroupFlag(parser):
  """Adds a --dump-path or --dump-flags flag to the given parser."""
  dump_group = parser.add_group(mutex=True)
  AddDumpPathFlag(dump_group)
  AddDumpFlagsFlag(dump_group)


def AddDumpPathFlag(parser):
  """Adds a --dump-path flag to the given parser."""
  help_text = """\
    Path to the dump file in Google Cloud Storage, in the format:
    `gs://[BUCKET_NAME]/[OBJECT_NAME]`.
    """
  parser.add_argument('--dump-path', help=help_text)


def AddDumpFlagsFlag(parser):
  """Adds a --dump-flags flag to the given parser."""
  help_text = """\
    A list of dump flags. An object containing a list of "key": "value" pairs.
    """
  parser.add_argument(
      '--dump-flags',
      metavar='KEY=VALUE',
      type=arg_parsers.ArgDict(),
      help=help_text,
  )


def AddConnectivityGroupFlag(parser, api_type, required=False):
  """Adds connectivity flag group to the given parser."""
  if api_type == ApiType.CREATE:
    connectivity_group = parser.add_group(
        (
            'The connectivity method used by the migration job. If a'
            " connectivity method isn't specified, then it isn't added to the"
            ' migration job.'
        ),
        mutex=True,
    )
  elif api_type == ApiType.UPDATE:
    connectivity_group = parser.add_group(
        (
            'The connectivity method used by the migration job. If a'
            " connectivity method isn't specified, then it isn't updated for"
            ' the migration job.'
        ),
        mutex=True,
    )
  connectivity_group.add_argument(
      '--static-ip',
      action='store_true',
      help=(
          'Use the default IP allowlist method. This method creates a public IP'
          ' that will be used with the destination Cloud SQL database. The'
          ' method works by configuring the source database server to accept'
          ' connections from the outgoing IP of the Cloud SQL instance.'
      ),
  )
  connectivity_group.add_argument(
      '--peer-vpc',
      help=(
          'Name of the VPC network to peer with the Cloud SQL private network.'
      ),
  )
  reverse_ssh_group = connectivity_group.add_group(
      'Parameters for the reverse-SSH tunnel connectivity method.'
  )
  reverse_ssh_group.add_argument(
      '--vm-ip', help='Bastion Virtual Machine IP.', required=required
  )
  reverse_ssh_group.add_argument(
      '--vm-port',
      help='Forwarding port for the SSH tunnel.',
      type=int,
      required=required,
  )
  reverse_ssh_group.add_argument(
      '--vm', help='Name of VM that will host the SSH tunnel bastion.'
  )
  reverse_ssh_group.add_argument(
      '--vpc',
      help='Name of the VPC network where the VM is hosted.',
      required=required,
  )


def AddFilterFlag(parser):
  """Adds a --filter flag to the given parser."""
  help_text = (
      'Filter the entities based on [AIP-160](https://google.aip.dev/160)'
      ' standard. Example: to filter all tables whose name start with'
      ' "Employee" and are present under schema "Company", use filter as'
      ' "Company.Employee```*``` AND type=TABLE"'
  )
  parser.add_argument('--filter', help=help_text)


def AddCommitIdFlag(parser):
  """Adds a --commit-id flag to the given parser."""
  help_text = (
      'Commit id for the conversion workspace to use for creating the migration'
      ' job. If not specified, the latest commit id will be used by default.'
  )
  parser.add_argument('--commit-id', help=help_text)


def AddDumpParallelLevelFlag(parser):
  """Adds a --dump-parallel-level flag to the given parser."""
  help_text = (
      'Parallelization level during initial dump of the migration job. If not'
      ' specified, will be defaulted to OPTIMAL.'
  )
  choices = ['MIN', 'OPTIMAL', 'MAX']
  parser.add_argument('--dump-parallel-level', help=help_text, choices=choices)


def AddDumpTypeFlag(parser):
  """Adds a --dump-type flag to the given parser."""
  help_text = (
      'The type of the data dump. Currently applicable for MySQL to MySQL '
      'migrations only.'
  )
  choices = ['LOGICAL', 'PHYSICAL']
  parser.add_argument('--dump-type', help=help_text, choices=choices)


def AddSqlServerHomogeneousMigrationConfigFlag(parser, is_update=False):
  """Adds SQL Server homogeneous migration flag group to the given parser."""
  sqlserver_homogeneous_migration_config = parser.add_group(mutex=True)
  AddSqlServerLogShippingConfigFlag(
      sqlserver_homogeneous_migration_config, is_update
  )
  AddSqlServerDagConfigFlag(sqlserver_homogeneous_migration_config, is_update)


def AddSqlServerLogShippingConfigFlag(parser, is_update=False):
  """Adds SQL Server homogeneous log shipping migration flag group to the given parser."""
  sqlserver_log_shipping_config = parser.add_group(
      (
          'The SQL Server homogeneous migration config. This is used only for'
          ' SQL Server to CloudSQL SQL Server migrations.'
      ),
  )
  if is_update:
    AddSqlServerBackupFilePattern(
        sqlserver_log_shipping_config, default_value=None
    )
    AddSqlServerDatabasesFlag(
        sqlserver_log_shipping_config, required=False
    )
    AddSqlServerEncryptedDatabasesFlag(sqlserver_log_shipping_config)
    AddSqlServerUseDiffBackupFlag(sqlserver_log_shipping_config)
    AddSqlServerPromoteWhenReadyFlag(sqlserver_log_shipping_config)
  else:
    AddSqlServerBackupFilePattern(
        sqlserver_log_shipping_config,
        default_value='.*(\\.|_)(<epoch>\\d*)\\.(trn|bak|trn\\.final)',
    )
    AddSqlServerDatabasesFlag(
        sqlserver_log_shipping_config, required=True
    )
    AddSqlServerEncryptedDatabasesFlag(sqlserver_log_shipping_config)
    AddSqlServerUseDiffBackupFlag(sqlserver_log_shipping_config)
    AddSqlServerPromoteWhenReadyFlag(sqlserver_log_shipping_config)


def AddSqlServerBackupFilePattern(parser, default_value=None):
  """Adds a --sqlserver-backup-file-pattern flag to the given parser."""
  help_text = (
      'Pattern that describes the default backup naming strategy. The specified'
      ' pattern should ensure lexicographical order of backups. The pattern'
      ' should define the following capture group set\nepoch - unix'
      ' timestamp\nExample: For backup files TestDB.1691448240.bak,'
      ' TestDB.1691448254.trn, TestDB.1691448284.trn.final use pattern:'
      ' .*\\.(<epoch>\\d*)\\.(trn|bak|trn\\.final) or'
      ' .*\\.(<timestamp>\\d*)\\.(trn|bak|trn\\.final)\nThis flag is used only'
      ' for SQL Server to Cloud SQL migrations.'
  )
  parser.add_argument(
      '--sqlserver-backup-file-pattern',
      help=help_text,
      default=default_value,
      hidden=True,
  )


def AddSqlServerDatabasesFlag(parser, required=True):
  """Adds a --sqlserver-databases flag to the given parser."""
  help_text = """\
    A list of databases to be migrated to the destination Cloud SQL instance.
    Provide databases as a comma separated list. This list should contain all
    encrypted and non-encrypted database names. This flag is used only for
    SQL Server to Cloud SQL migrations.
    """
  parser.add_argument(
      '--sqlserver-databases',
      metavar='databaseName',
      type=arg_parsers.ArgList(min_length=1),
      help=help_text,
      required=required,
  )


def AddSqlServerEncryptedDatabasesFlag(parser):
  """Adds a --sqlserver-encrypted-databases flag to the given parser."""
  help_text = """\
    A JSON/YAML file describing the encryption settings per database for all encrytped databases.
    Note:
    Path to the Certificate (.cer) and Private Key (.pvk) in Cloud Storage,
    should be in the form of `gs://bucketName/fileName`. The instance must
    have write permissions to the bucket and read access to the file.
    An example of a JSON request:
        [{
            "database": "db1",
            "encryptionOptions": {
                "certPath": "Path to certificate 1",
                "pvkPath": "Path to certificate private key 1",
                "pvkPassword": "Private key password 1"
            }
        },
        {
            "database": "db2",
            "encryptionOptions": {
                "certPath": "Path to certificate 2",
                "pvkPath": "Path to certificate private key 2",
                "pvkPassword": "Private key password 2"
            }
        }]

      This flag accepts "-" for stdin. This flag is used only for SQL Server to Cloud SQL migrations.
    """
  parser.add_argument(
      '--sqlserver-encrypted-databases',
      type=arg_parsers.YAMLFileContents(),
      help=help_text,
  )


def AddSqlServerUseDiffBackupFlag(parser):
  """Adds a --sqlserver-diff-backup flag to the given parser."""
  help_text = """\
      Enable differential backups. If not specified, differential backups
      are disabled by default. Use --sqlserver-diff-backup to enable and
      --no-sqlserver-diff-backup to disable. This flag is used only for
      homogeneous SQL Server to Cloud SQL for SQL Server migrations.
    """
  parser.add_argument(
      '--sqlserver-diff-backup',
      action='store_true',
      help=help_text,
  )


def AddSqlServerDagConfigFlag(parser, is_update=False):
  """Adds SQL Server homogeneous DAG migration flag group to the given parser."""
  sqlserver_dag_config = parser.add_group(
      (
          'The SQL Server homogeneous DAG migration config. This is used only'
          ' for SQL Server to CloudSQL SQL for SQL Server migrations.'
      ),
      hidden=True,
  )
  sqlserver_dag_config.add_argument(
      '--sqlserver-dag-source-ag',
      help='The name of the source availability group.',
      required=not is_update,
      hidden=True,
  )
  sqlserver_dag_config.add_argument(
      '--sqlserver-dag-linked-server',
      help=(
          'The name of the linked server that points to the source SQL Server'
          ' instance.'
      ),
      required=not is_update,
      hidden=True,
  )


def AddSkipValidationFlag(parser):
  """Adds a --skip-validation flag to the given parser."""
  help_text = """\
    Restart the migration job without running prior configuration verification.
    """
  parser.add_argument(
      '--skip-validation',
      action='store_true',
      help=help_text,
  )


def AddMigrationJobObjectsConfigFlagForCreateAndUpdate(parser):
  """Adds migration job objects config flag group to the given parser."""
  migration_config = parser.add_group(
      'The migration job objects config.',
      mutex=True,
  )
  database_config = migration_config.add_group(
      'The migration job objects config for databases.',
      mutex=True,
  )
  AddDatabasesFilterFlag(database_config)
  AddAllDatabasesFlag(database_config)


def AddMigrationJobObjectsConfigFlagForPromote(parser):
  """Adds migration job objects config flag group to the given parser."""
  migration_config = parser.add_group(
      'The migration job objects config.',
      mutex=True,
  )
  AddDatabasesFilterFlagForSqlServer(migration_config)


def AddMigrationJobObjectsConfigFlagForRestart(parser):
  """Adds migration job objects config flag group to the given parser."""
  migration_config = parser.add_group(
      'The migration job objects config.',
      mutex=True,
  )
  AddDatabasesFilterFlagForSqlServer(migration_config)
  AddObjectFilterFlagForHeterogeneous(migration_config)


def AddDatabasesFilterFlag(parser):
  """Adds a --databases-filter flag to the given parser."""
  help_text = """\
    A list of databases to be migrated to the destination instance.
    Provide databases as a comma separated list. This flag is used only for
    Postgres to AlloyDB, and Postgres to Cloud SQL Postgres and MySQL to Cloud
    SQL MySQL migrations.
    """
  parser.add_argument(
      '--databases-filter',
      metavar='databaseName',
      type=arg_parsers.ArgList(min_length=1),
      help=help_text,
  )


def AddAllDatabasesFlag(parser):
  """Adds --all-databases flag to the given parser."""
  help_text = """\
    Migrate all databases for the migration job. This flag is used only for
    Postgres to AlloyDB, and Postgres to Cloud SQL Postgres and MySQL to Cloud
    SQL MySQL migrations.
    """
  parser.add_argument('--all-databases', action='store_true', help=help_text)


def AddDatabasesFilterFlagForSqlServer(parser):
  """Adds a --databases-filter flag to the given parser."""
  help_text = """\
    A list of databases to be migrated to the destination instance.
    Provide databases as a comma separated list. This flag is used only for
    SQL Server to Cloud SQL SQL Server migrations.
    """
  parser.add_argument(
      '--databases-filter',
      metavar='databaseName',
      type=arg_parsers.ArgList(min_length=1),
      help=help_text,
  )


def AddObjectFilterFlagForHeterogeneous(parser):
  """Adds a --object-filter flag to the given parser."""
  help_text = """\
    A list of schema and table names to be migrated to the destination instance.
    Usage: --object-filter schema=schema1,table=table1 --object-filter schema=schema2,table=table2
    This flag is used only for heterogeneous migrations.
    """
  parser.add_argument(
      '--object-filter',
      type=arg_parsers.ArgDict(
          spec={
              'schema': str,
              'table': str,
          },
          required_keys=['schema', 'table'],
      ),
      action='append',
      help=help_text,
      hidden=True,
  )


def AddSqlServerPromoteWhenReadyFlag(parser):
  """Adds a --sqlserver-promote-when-ready flag to the given parser."""
  help_text = """\
      Promote the database when it is ready. Use --sqlserver-promote-when-ready
      to enable and --no-sqlserver-promote-when-ready to disable. This flag is
      used only for homogeneous SQL Server to Cloud SQL for SQL Server
      migrations.
    """
  parser.add_argument(
      '--sqlserver-promote-when-ready',
      action='store_true',
      help=help_text,
  )


def AddRestartFailedObjectsFlag(parser):
  """Adds a --restart-failed-objects flag to the given parser."""
  help_text = """\
    Restart the failed objects in the migration job. This flag is used only for
    Postgres to AlloyDB, and Postgres to Cloud SQL Postgres and MySQL to Cloud
    SQL MySQL migrations.
    """
  parser.add_argument(
      '--restart-failed-objects',
      action='store_true',
      help=help_text,
  )


def AddHeterogeneousMigrationConfigFlag(parser, is_update=False):
  """Adds heterogeneous migration flag group to the given parser."""
  heterogeneous_migration_config = parser.add_group(
      (
          'The heterogeneous migration config. This is used only for'
          ' Oracle to Cloud SQL for PostgreSQL and SQL Server to Cloud SQL for'
          ' PostgreSQL migrations.'
      ),
  )
  AddHeterogeneousMigrationSourceConfigFlags(
      heterogeneous_migration_config, is_update
  )
  AddHeterogeneousMigrationDestinationConfigFlags(
      heterogeneous_migration_config
  )


def AddHeterogeneousMigrationDestinationConfigFlags(parser):
  """Adds heterogeneous migration destination config flag to the parser."""
  destination_config = parser.add_group(
      (
          'Configuration for Postgres as a destination in a heterogeneous'
          ' migration.'
      ),
  )
  destination_config.add_argument(
      '--max-concurrent-destination-connections',
      help="""\
        Maximum number of concurrent connections Database Migration Service will
        open to the destination for data migration.
        """,
      type=int,
  )
  destination_config.add_argument(
      '--transaction-timeout',
      help="""Timeout for data migration transactions.""",
      type=arg_parsers.Duration(lower_bound='30s', upper_bound='300s'),
  )


def AddHeterogeneousMigrationSourceConfigFlags(parser, is_update=False):
  """Adds heterogeneous migration source config flag group to the parser."""
  source_config = parser.add_group(
      (
          'Configuration for Oracle or SQL Server as a source in a'
          ' heterogeneous migration.'
      ),
  )
  source_config.add_argument(
      '--max-concurrent-full-dump-connections',
      help="""\
        Maximum number of connections Database Migration Service will open to
        the source for full dump phase.
        """,
      type=int,
  )
  source_config.add_argument(
      '--max-concurrent-cdc-connections',
      help="""\
        Maximum number of connections Database Migration Service will open to
        the source for CDC phase.
        """,
      type=int,
  )
  if not is_update:
    skip_full_dump_group = source_config.add_group(
        'Configuration for skipping full dump.',
    )
    skip_full_dump_group.add_argument(
        '--skip-full-dump',
        help="""\
          Whether to skip full dump or not.
          """,
        action='store_true',
    )
    cdc_start_position_group = skip_full_dump_group.add_group(
        'Configuration for CDC start position.',
        mutex=True,
    )
    cdc_start_position_group.add_argument(
        '--oracle-cdc-start-position',
        help="""\
          Oracle schema change number (SCN) to start CDC data migration from.
          """,
        type=int,
    )
    cdc_start_position_group.add_argument(
        '--sqlserver-cdc-start-position',
        help="""\
          Sqlserver log squence number (LSN) to start CDC data migration from.
          """
    )


def AddPostgresHomogeneousConfigFlag(parser):
  """Adds postgres homogeneous config flag to the given parser."""
  parser.add_argument(
      '--use-postgres-native',
      action='store_true',
      default=False,
      help='Use native logical replication for PostgreSQL migration.',
      hidden=True,
  )