#!/usr/bin/env python3
# Copyright 2017 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 tests using docker images in Google Container Registry per matrix."""

import argparse
import atexit
import json
import multiprocessing
import os
import re
import subprocess
import sys
import uuid

# Language Runtime Matrix
import client_matrix

python_util_dir = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "../run_tests/python_utils")
)
sys.path.append(python_util_dir)
import dockerjob
import jobset
import report_utils
import upload_test_results

_TEST_TIMEOUT_SECONDS = 60
_PULL_IMAGE_TIMEOUT_SECONDS = 15 * 60
_MAX_PARALLEL_DOWNLOADS = 6
_LANGUAGES = list(client_matrix.LANG_RUNTIME_MATRIX.keys())
# All gRPC release tags, flattened, deduped and sorted.
_RELEASES = sorted(
    list(
        set(
            release
            for release_dict in list(client_matrix.LANG_RELEASE_MATRIX.values())
            for release in list(release_dict.keys())
        )
    )
)

argp = argparse.ArgumentParser(description="Run interop tests.")
argp.add_argument("-j", "--jobs", default=multiprocessing.cpu_count(), type=int)
argp.add_argument(
    "--docker_path",
    default="us-docker.pkg.dev/grpc-testing/testing-images-public",
    help="Path of docker images",
)
argp.add_argument(
    "--release",
    default="all",
    choices=["all"] + _RELEASES,
    help=(
        "Release tags to test.  When testing all "
        'releases defined in client_matrix.py, use "all".'
    ),
)
argp.add_argument(
    "-l",
    "--language",
    choices=["all"] + sorted(_LANGUAGES),
    nargs="+",
    default=["all"],
    help="Languages to test",
)
argp.add_argument(
    "--keep",
    action="store_true",
    help="keep the created local images after finishing the tests.",
)
argp.add_argument(
    "--report_file", default="report.xml", help="The result file to create."
)
argp.add_argument(
    "--allow_flakes",
    default=False,
    action="store_const",
    const=True,
    help=(
        "Allow flaky tests to show as passing (re-runs failed "
        "tests up to five times)"
    ),
)
argp.add_argument(
    "--bq_result_table",
    default="",
    type=str,
    nargs="?",
    help="Upload test results to a specified BQ table.",
)
# Requests will be routed through specified VIP by default.
# See go/grpc-interop-tests (internal-only) for details.
argp.add_argument(
    "--server_host",
    default="74.125.206.210",
    type=str,
    nargs="?",
    help="The gateway to backend services.",
)


def _get_test_images_for_lang(lang, release_arg, image_path_prefix):
    """Find docker images for a language across releases and runtimes.

    Returns dictionary of list of (<tag>, <image-full-path>) keyed by runtime.
    """
    if release_arg == "all":
        # Use all defined releases for given language
        releases = client_matrix.get_release_tags(lang)
    else:
        # Look for a particular release.
        if release_arg not in client_matrix.get_release_tags(lang):
            jobset.message(
                "SKIPPED",
                "release %s for %s is not defined" % (release_arg, lang),
                do_newline=True,
            )
            return {}
        releases = [release_arg]

    # Image tuples keyed by runtime.
    images = {}
    for tag in releases:
        for runtime in client_matrix.get_runtimes_for_lang_release(lang, tag):
            image_name = "%s/grpc_interop_%s:%s" % (
                image_path_prefix,
                runtime,
                tag,
            )
            image_tuple = (tag, image_name)

            if runtime not in images:
                images[runtime] = []
            images[runtime].append(image_tuple)
    return images


def _read_test_cases_file(lang, runtime, release):
    """Read test cases from a bash-like file and return a list of commands"""
    # Check to see if we need to use a particular version of test cases.
    release_info = client_matrix.LANG_RELEASE_MATRIX[lang].get(release)
    if release_info:
        testcases_file = release_info.testcases_file
    if not testcases_file:
        # TODO(jtattermusch): remove the double-underscore, it is pointless
        testcases_file = "%s__master" % lang

    # For csharp, the testcases file used depends on the runtime
    # TODO(jtattermusch): remove this odd special case
    if lang == "csharp" and runtime == "csharpcoreclr":
        testcases_file = testcases_file.replace("csharp_", "csharpcoreclr_")

    testcases_filepath = os.path.join(
        os.path.dirname(__file__), "testcases", testcases_file
    )
    lines = []
    with open(testcases_filepath) as f:
        for line in f.readlines():
            line = re.sub("\\#.*$", "", line)  # remove hash comments
            line = line.strip()
            if line and not line.startswith("echo"):
                # Each non-empty line is a treated as a test case command
                lines.append(line)
    return lines


def _cleanup_docker_image(image):
    jobset.message("START", "Cleanup docker image %s" % image, do_newline=True)
    dockerjob.remove_image(image, skip_nonexistent=True)


args = argp.parse_args()


# caches test cases (list of JobSpec) loaded from file.  Keyed by lang and runtime.
def _generate_test_case_jobspecs(lang, runtime, release, suite_name):
    """Returns the list of test cases from testcase files per lang/release."""
    testcase_lines = _read_test_cases_file(lang, runtime, release)

    job_spec_list = []
    for line in testcase_lines:
        print("Creating jobspec with cmdline '{}'".format(line))
        # TODO(jtattermusch): revisit the logic for updating test case commands
        # what it currently being done seems fragile.

        # Extract test case name from the command line
        m = re.search(r"--test_case=(\w+)", line)
        testcase_name = m.group(1) if m else "unknown_test"

        # Extract the server name from the command line
        if "--server_host_override=" in line:
            m = re.search(
                r"--server_host_override=((.*).sandbox.googleapis.com)", line
            )
        else:
            m = re.search(r"--server_host=((.*).sandbox.googleapis.com)", line)
        server = m.group(1) if m else "unknown_server"
        server_short = m.group(2) if m else "unknown_server"

        # replace original server_host argument
        assert "--server_host=" in line
        line = re.sub(
            r"--server_host=[^ ]*", r"--server_host=%s" % args.server_host, line
        )

        # some interop tests don't set server_host_override (see #17407),
        # but we need to use it if different host is set via cmdline args.
        if args.server_host != server and not "--server_host_override=" in line:
            line = re.sub(
                r"(--server_host=[^ ]*)",
                r"\1 --server_host_override=%s" % server,
                line,
            )

        spec = jobset.JobSpec(
            cmdline=line,
            shortname="%s:%s:%s:%s"
            % (suite_name, lang, server_short, testcase_name),
            timeout_seconds=_TEST_TIMEOUT_SECONDS,
            shell=True,
            flake_retries=5 if args.allow_flakes else 0,
        )
        job_spec_list.append(spec)
    return job_spec_list


def _pull_image_for_lang(lang, image, release):
    """Pull an image for a given language form the image registry."""
    cmdline = [
        "time gcloud docker -- pull %s && time docker run --rm=true %s"
        " /bin/true" % (image, image)
    ]
    return jobset.JobSpec(
        cmdline=cmdline,
        shortname="pull_image_{}".format(image),
        timeout_seconds=_PULL_IMAGE_TIMEOUT_SECONDS,
        shell=True,
        flake_retries=2,
    )


def _test_release(lang, runtime, release, image, xml_report_tree, skip_tests):
    total_num_failures = 0
    suite_name = "%s__%s_%s" % (lang, runtime, release)
    job_spec_list = _generate_test_case_jobspecs(
        lang, runtime, release, suite_name
    )

    if not job_spec_list:
        jobset.message("FAILED", "No test cases were found.", do_newline=True)
        total_num_failures += 1
    else:
        num_failures, resultset = jobset.run(
            job_spec_list,
            newline_on_success=True,
            add_env={"docker_image": image},
            maxjobs=args.jobs,
            skip_jobs=skip_tests,
        )
        if args.bq_result_table and resultset:
            upload_test_results.upload_interop_results_to_bq(
                resultset, args.bq_result_table
            )
        if skip_tests:
            jobset.message("FAILED", "Tests were skipped", do_newline=True)
            total_num_failures += 1
        if num_failures:
            total_num_failures += num_failures

        report_utils.append_junit_xml_results(
            xml_report_tree,
            resultset,
            "grpc_interop_matrix",
            suite_name,
            str(uuid.uuid4()),
        )
    return total_num_failures


def _run_tests_for_lang(lang, runtime, images, xml_report_tree):
    """Find and run all test cases for a language.

    images is a list of (<release-tag>, <image-full-path>) tuple.
    """
    skip_tests = False
    total_num_failures = 0

    max_pull_jobs = min(args.jobs, _MAX_PARALLEL_DOWNLOADS)
    max_chunk_size = max_pull_jobs
    chunk_count = (len(images) + max_chunk_size) // max_chunk_size

    for chunk_index in range(chunk_count):
        chunk_start = chunk_index * max_chunk_size
        chunk_size = min(max_chunk_size, len(images) - chunk_start)
        chunk_end = chunk_start + chunk_size
        pull_specs = []
        if not skip_tests:
            for release, image in images[chunk_start:chunk_end]:
                pull_specs.append(_pull_image_for_lang(lang, image, release))

        # NOTE(rbellevi): We batch docker pull operations to maximize
        # parallelism, without letting the disk usage grow unbounded.
        pull_failures, _ = jobset.run(
            pull_specs, newline_on_success=True, maxjobs=max_pull_jobs
        )
        if pull_failures:
            jobset.message(
                "FAILED",
                'Image download failed. Skipping tests for language "%s"'
                % lang,
                do_newline=True,
            )
            skip_tests = True
        for release, image in images[chunk_start:chunk_end]:
            total_num_failures += _test_release(
                lang, runtime, release, image, xml_report_tree, skip_tests
            )
        if not args.keep:
            for _, image in images[chunk_start:chunk_end]:
                _cleanup_docker_image(image)
    if not total_num_failures:
        jobset.message(
            "SUCCESS", "All {} tests passed".format(lang), do_newline=True
        )
    else:
        jobset.message(
            "FAILED", "Some {} tests failed".format(lang), do_newline=True
        )

    return total_num_failures


languages = args.language if args.language != ["all"] else _LANGUAGES
total_num_failures = 0
_xml_report_tree = report_utils.new_junit_xml_tree()
for lang in languages:
    docker_images = _get_test_images_for_lang(
        lang, args.release, args.docker_path
    )
    for runtime in sorted(docker_images.keys()):
        total_num_failures += _run_tests_for_lang(
            lang, runtime, docker_images[runtime], _xml_report_tree
        )

report_utils.create_xml_report_file(_xml_report_tree, args.report_file)

if total_num_failures:
    sys.exit(1)
sys.exit(0)
