#!/usr/bin/env python
# Copyright 2020 gRPC authors.
#
# 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.
"""Run xDS integration tests on GCP using Traffic Director."""

import argparse
import googleapiclient.discovery
import grpc
import logging
import os
import random
import shlex
import socket
import subprocess
import sys
import tempfile
import time

from oauth2client.client import GoogleCredentials

import python_utils.jobset as jobset
import python_utils.report_utils as report_utils

from src.proto.grpc.testing import empty_pb2
from src.proto.grpc.testing import messages_pb2
from src.proto.grpc.testing import test_pb2_grpc

logger = logging.getLogger()
console_handler = logging.StreamHandler()
formatter = logging.Formatter(fmt='%(asctime)s: %(levelname)-8s %(message)s')
console_handler.setFormatter(formatter)
logger.handlers = []
logger.addHandler(console_handler)
logger.setLevel(logging.WARNING)

_TEST_CASES = [
    'backends_restart',
    'change_backend_service',
    'gentle_failover',
    'new_instance_group_receives_traffic',
    'ping_pong',
    'remove_instance_group',
    'round_robin',
    'secondary_locality_gets_no_requests_on_partial_primary_failure',
    'secondary_locality_gets_requests_on_primary_failure',
    'traffic_splitting',
]
# Valid test cases, but not in all. So the tests can only run manually, and
# aren't enabled automatically for all languages.
#
# TODO: Move them into _TEST_CASES when support is ready in all languages.
_ADDITIONAL_TEST_CASES = ['path_matching', 'header_matching']


def parse_test_cases(arg):
    if arg == '':
        return []
    arg_split = arg.split(',')
    test_cases = set()
    all_test_cases = _TEST_CASES + _ADDITIONAL_TEST_CASES
    for arg in arg_split:
        if arg == "all":
            test_cases = test_cases.union(_TEST_CASES)
        else:
            test_cases = test_cases.union([arg])
    if not all([test_case in all_test_cases for test_case in test_cases]):
        raise Exception('Failed to parse test cases %s' % arg)
    # Perserve order.
    return [x for x in all_test_cases if x in test_cases]


def parse_port_range(port_arg):
    try:
        port = int(port_arg)
        return range(port, port + 1)
    except:
        port_min, port_max = port_arg.split(':')
        return range(int(port_min), int(port_max) + 1)


argp = argparse.ArgumentParser(description='Run xDS interop tests on GCP')
argp.add_argument('--project_id', help='GCP project id')
argp.add_argument(
    '--gcp_suffix',
    default='',
    help='Optional suffix for all generated GCP resource names. Useful to '
    'ensure distinct names across test runs.')
argp.add_argument(
    '--test_case',
    default='ping_pong',
    type=parse_test_cases,
    help='Comma-separated list of test cases to run. Available tests: %s, '
    '(or \'all\' to run every test). '
    'Alternative tests not included in \'all\': %s' %
    (','.join(_TEST_CASES), ','.join(_ADDITIONAL_TEST_CASES)))
argp.add_argument(
    '--bootstrap_file',
    default='',
    help='File to reference via GRPC_XDS_BOOTSTRAP. Disables built-in '
    'bootstrap generation')
argp.add_argument(
    '--client_cmd',
    default=None,
    help='Command to launch xDS test client. {server_uri}, {stats_port} and '
    '{qps} references will be replaced using str.format(). GRPC_XDS_BOOTSTRAP '
    'will be set for the command')
argp.add_argument('--zone', default='us-central1-a')
argp.add_argument('--secondary_zone',
                  default='us-west1-b',
                  help='Zone to use for secondary TD locality tests')
argp.add_argument('--qps', default=100, type=int, help='Client QPS')
argp.add_argument(
    '--wait_for_backend_sec',
    default=1200,
    type=int,
    help='Time limit for waiting for created backend services to report '
    'healthy when launching or updated GCP resources')
argp.add_argument(
    '--use_existing_gcp_resources',
    default=False,
    action='store_true',
    help=
    'If set, find and use already created GCP resources instead of creating new'
    ' ones.')
argp.add_argument(
    '--keep_gcp_resources',
    default=False,
    action='store_true',
    help=
    'Leave GCP VMs and configuration running after test. Default behavior is '
    'to delete when tests complete.')
argp.add_argument(
    '--compute_discovery_document',
    default=None,
    type=str,
    help=
    'If provided, uses this file instead of retrieving via the GCP discovery '
    'API')
argp.add_argument(
    '--alpha_compute_discovery_document',
    default=None,
    type=str,
    help='If provided, uses this file instead of retrieving via the alpha GCP '
    'discovery API')
argp.add_argument('--network',
                  default='global/networks/default',
                  help='GCP network to use')
argp.add_argument('--service_port_range',
                  default='8080:8110',
                  type=parse_port_range,
                  help='Listening port for created gRPC backends. Specified as '
                  'either a single int or as a range in the format min:max, in '
                  'which case an available port p will be chosen s.t. min <= p '
                  '<= max')
argp.add_argument(
    '--stats_port',
    default=8079,
    type=int,
    help='Local port for the client process to expose the LB stats service')
argp.add_argument('--xds_server',
                  default='trafficdirector.googleapis.com:443',
                  help='xDS server')
argp.add_argument('--source_image',
                  default='projects/debian-cloud/global/images/family/debian-9',
                  help='Source image for VMs created during the test')
argp.add_argument('--path_to_server_binary',
                  default=None,
                  type=str,
                  help='If set, the server binary must already be pre-built on '
                  'the specified source image')
argp.add_argument('--machine_type',
                  default='e2-standard-2',
                  help='Machine type for VMs created during the test')
argp.add_argument(
    '--instance_group_size',
    default=2,
    type=int,
    help='Number of VMs to create per instance group. Certain test cases (e.g., '