#!/usr/bin/env python3
# Copyright 2022 The 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.
"""Helps with running bazel with extra settings to generate structured test reports in CI."""

import argparse
import os
import platform
import sys
import uuid

_ROOT = os.path.abspath(os.path.join(os.path.dirname(sys.argv[0]), '../../..'))
os.chdir(_ROOT)

# How long to sleep before querying Resultstore API and uploading to bigquery
# (to let ResultStore finish writing results from the bazel invocation that has
# just finished).
_UPLOAD_RBE_RESULTS_DELAY_SECONDS = 60


def _platform_string():
    """Detect current platform"""
    if platform.system() == 'Windows':
        return 'windows'
    elif platform.system()[:7] == 'MSYS_NT':
        return 'windows'
    elif platform.system() == 'Darwin':
        return 'mac'
    elif platform.system() == 'Linux':
        return 'linux'
    else:
        return 'posix'


def _append_to_kokoro_bazel_invocations(invocation_id: str) -> None:
    """Kokoro can display "Bazel" result link on kokoro jobs if told so."""
    # to get "bazel" link for kokoro build, we need to upload
    # the "bazel_invocation_ids" file with bazel invocation ID as artifact.
    kokoro_artifacts_dir = os.getenv('KOKORO_ARTIFACTS_DIR')
    if kokoro_artifacts_dir:
        # append the bazel invocation UUID to the bazel_invocation_ids file.
        with open(os.path.join(kokoro_artifacts_dir, 'bazel_invocation_ids'),
                  'a') as f:
            f.write(invocation_id + '\n')
        print(
            'Added invocation ID %s to kokoro "bazel_invocation_ids" artifact' %
            invocation_id,
            file=sys.stderr)
    else:
        print(
            'Skipped adding invocation ID %s to kokoro "bazel_invocation_ids" artifact'
            % invocation_id,
            file=sys.stderr)
        pass


def _generate_junit_report_string(report_suite_name: str, invocation_id: str,
                                  success: bool) -> None:
    """Generate sponge_log.xml formatted report, that will make the bazel invocation reachable as a target in resultstore UI / sponge."""
    bazel_invocation_url = 'https://source.cloud.google.com/results/invocations/%s' % invocation_id
    package_name = report_suite_name
    # set testcase name to invocation URL. That way, the link will be displayed in some form
    # resultstore UI and sponge even in case the bazel invocation succeeds.
    testcase_name = bazel_invocation_url
    if success:
        # unfortunately, neither resultstore UI nor sponge display the "system-err" output (or any other tags)
        # on a passing test case. But at least we tried.
        test_output_tag = '<system-err>PASSED. See invocation results here: %s</system-err>' % bazel_invocation_url
    else:
        # The failure output will be displayes in both resultstore UI and sponge when clicking on the failing testcase.
        test_output_tag = '<failure message="Failure">FAILED. See bazel invocation results here: %s</failure>' % bazel_invocation_url

    lines = [
        '<testsuites>',
        '<testsuite id="1" name="%s" package="%s">' %
        (report_suite_name, package_name),
        '<testcase name="%s">' % testcase_name,
        test_output_tag,
        '</testcase>'
        '</testsuite>',
        '</testsuites>',
    ]
    return '\n'.join(lines)


def _create_bazel_wrapper(report_path: str, report_suite_name: str,
                          invocation_id: str, upload_results: bool) -> None:
    """Create a "bazel wrapper" script that will execute bazel with extra settings and postprocessing."""

    os.makedirs(report_path, exist_ok=True)

    bazel_wrapper_filename = os.path.join(report_path, 'bazel_wrapper')
    bazel_wrapper_bat_filename = bazel_wrapper_filename + '.bat'
    bazel_rc_filename = os.path.join(report_path, 'bazel_wrapper.bazelrc')

    # put xml reports in a separate directory if requested by GRPC_TEST_REPORT_BASE_DIR
    report_base_dir = os.getenv('GRPC_TEST_REPORT_BASE_DIR', None)
    xml_report_path = os.path.abspath(
        os.path.join(report_base_dir, report_path
                    ) if report_base_dir else report_path)
    os.makedirs(xml_report_path, exist_ok=True)

    failing_report_filename = os.path.join(xml_report_path, 'sponge_log.xml')
    success_report_filename = os.path.join(xml_report_path,
                                           'success_log_to_rename.xml')

    if _platform_string() == 'windows':
        workspace_status_command = 'tools/remote_build/workspace_status_kokoro.bat'
    else:
        workspace_status_command = 'tools/remote_build/workspace_status_kokoro.sh'

    # generate RC file with the bazel flags we want to use apply.
    # Using an RC file solves problems with flag ordering in the wrapper.
    # (e.g. some flags need to come after the build/test command)
    with open(bazel_rc_filename, 'w') as f:
        f.write('build --invocation_id="%s"\n' % invocation_id)
        f.write('build --workspace_status_command="%s"\n' %
                workspace_status_command)

    # generate "failing" and "success" report
    # the "failing" is named as "sponge_log.xml", which is the name picked up by sponge/resultstore
    # so the failing report will be used by default (unless we later replace the report with
    # one that says "success"). That way if something goes wrong before bazel is run,
    # there will at least be a "failing" target that indicates that (we really don't want silent failures).
    with open(failing_report_filename, 'w') as f:
        f.write(
            _generate_junit_report_string(report_suite_name,
                                          invocation_id,
                                          success=False))
    with open(success_report_filename, 'w') as f:
        f.write(
            _generate_junit_report_string(report_suite_name,
                                          invocation_id,
                                          success=True))

    # generate the bazel wrapper for linux/macos
    with open(bazel_wrapper_filename, 'w') as f:
        intro_lines = [
            '#!/bin/bash',
            'set -ex',
            '',
            'tools/bazel --bazelrc="%s" "$@" || FAILED=true' %
            bazel_rc_filename,
            '',
        ]

        if upload_results:
            upload_results_lines = [
                'sleep %s' % _UPLOAD_RBE_RESULTS_DELAY_SECONDS,
                'PYTHONHTTPSVERIFY=0 python3 ./tools/run_tests/python_utils/upload_rbe_results.py --invocation_id="%s"'
                % invocation_id,
                '',
            ]
        else:
            upload_results_lines = []

        outro_lines = [
            'if [ "$FAILED" != "" ]',
            'then',
            '  exit 1',
            'else',
            '  # success: plant the pre-generated xml report that says "success"',
            '  mv -f %s %s' %
            (success_report_filename, failing_report_filename),
            'fi',
        ]

        lines = [
            line + '\n'
            for line in intro_lines + upload_results_lines + outro_lines
        ]
        f.writelines(lines)
    os.chmod(bazel_wrapper_filename, 0o775)  # make the unix wrapper executable

    # generate bazel wrapper for windows
    with open(bazel_wrapper_bat_filename, 'w') as f:
        intro_lines = [
            '@echo on',
            '',
            'bazel --bazelrc="%s" %%*' % bazel_rc_filename,
            'set BAZEL_EXITCODE=%errorlevel%',
            '',
        ]

        if upload_results:
            upload_results_lines = [
                'sleep %s' % _UPLOAD_RBE_RESULTS_DELAY_SECONDS,
                'python3 tools/run_tests/python_utils/upload_rbe_results.py --invocation_id="%s" || exit /b 1'
                % invocation_id,
                '',
            ]
        else:
            upload_results_lines = []

        outro_lines = [
            'if %BAZEL_EXITCODE% == 0 (',
            '  @rem success: plant the pre-generated xml report that says "success"',
            '  mv -f %s %s' %
            (success_report_filename, failing_report_filename),
            ')',
            'exit /b %BAZEL_EXITCODE%',
        ]

        lines = [
            line + '\n'
            for line in intro_lines + upload_results_lines + outro_lines
        ]
        f.writelines(lines)

    print('Bazel invocation ID: %s' % invocation_id, file=sys.stderr)
    print('Upload test results to BigQuery after bazel runs: %s' %
          upload_results,
          file=sys.stderr)
    print('Generated bazel wrapper: %s' % bazel_wrapper_filename,
          file=sys.stderr)
    print('Generated bazel wrapper: %s' % bazel_wrapper_bat_filename,
          file=sys.stderr)


if __name__ == '__main__':
    # parse command line
    argp = argparse.ArgumentParser(
        description=
        'Generate bazel wrapper to help with bazel test reports in CI.')
    argp.add_argument(
        '--report_path',
        required=True,
        type=str,
        help=
        'Path under which the bazel wrapper and other files are going to be generated'
    )
    argp.add_argument('--report_suite_name',
                      default='bazel_invocations',
                      type=str,
                      help='Test suite name to use in generated XML report')
    args = argp.parse_args()

    # generate new bazel invocation ID
    invocation_id = str(uuid.uuid4())

    report_path = args.report_path
    report_suite_name = args.report_suite_name
    upload_results = True if os.getenv('UPLOAD_TEST_RESULTS') else False

    _append_to_kokoro_bazel_invocations(invocation_id)
    _create_bazel_wrapper(report_path, report_suite_name, invocation_id,
                          upload_results)