#!/usr/bin/env python3
# ===--- run_cperf --------------------------------------------------------===
#
#  This source file is part of the Swift.org open source project
#
#  Copyright (c) 2014 - 2022 Apple Inc. and the Swift project authors
#  Licensed under Apache License v2.0 with Runtime Library Exception
#
#  See https://swift.org/LICENSE.txt for license information
#  See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
#
# ===----------------------------------------------------------------------===

"""A run script to be executed as a pull-request test on Jenkins."""

import argparse
import json
import os
import platform
import random
import re
import shutil
import sys
import unittest

import common


OLD_INSTANCE = 'old'
NEW_INSTANCE = 'new'


def get_workspace_for_instance(instance, args):
    if instance == OLD_INSTANCE and args.old_workspace is not None:
        return args.old_workspace
    elif instance == NEW_INSTANCE and args.new_workspace is not None:
        return args.new_workspace
    else:
        return common.private_workspace(
            os.path.join(args.swift_branch, instance))


def setup_workspace(instance, workspace, args):
    if args.clean and os.path.exists(workspace):
        shutil.rmtree(workspace)
    if not os.path.exists(workspace):
        os.makedirs(workspace)
    swift = os.path.join(workspace, "swift")
    if not os.path.exists(swift):
        common.git_clone(args.url, swift, tree=args.swift_branch)
    common.check_execute(['utils/update-checkout',
                          '--clone-with-ssh', '--reset-to-remote',
                          '--clean', '--scheme', args.swift_branch],
                         cwd=swift, timeout=9999999)
    if instance == NEW_INSTANCE:
        command_fetch = ['git', '-C', swift, 'fetch', 'origin',
                         'pull/%d/merge' % args.setup_workspaces_for_pr]
        common.check_execute(command_fetch)
        common.git_checkout('FETCH_HEAD', swift)


def validate_workspace(instance, workspace, args):
    if not os.path.exists(workspace):
        raise RuntimeError("Missing %s workspace dir %s"
                           % (instance, workspace))
    swift = os.path.join(workspace, "swift")
    if not os.path.exists(swift):
        raise RuntimeError("Missing swift checkout in workspace dir %s"
                           % workspace)


def workspace_comment():
    if 'WORKSPACE' in os.environ:
        workspace = os.environ['WORKSPACE']
        print("WORKSPACE: %s" % workspace)
        return os.path.join(workspace, 'comment.md')
    else:
        return None


# Utility class for testing run_cperf against various data discovered in
# the field; unfortunately we are not yet at the point where there's a
# real testsuite for it, but this is a step.
class SelfTest(unittest.TestCase):

    args = None

    def test_failed_modules(self):
        configs = get_configs(self.args.suite)
        (passed_modules, failed_modules) = find_module_statuses(configs,
                                                                self.args)
        print("passed: " + str(passed_modules))
        print("failed: " + str(failed_modules))


def main():

    args = parse_args()
    if args.self_test:
        SelfTest.args = args
        suite = unittest.TestLoader().loadTestsFromTestCase(SelfTest)
        return unittest.TextTestRunner(verbosity=2).run(suite)

    common.debug_print('** RUN PULL-REQUEST CPERF **')
    os.chdir(os.path.dirname(__file__))

    instances = [NEW_INSTANCE, OLD_INSTANCE]
    configs = get_configs(args.suite)

    ws_comment = workspace_comment()
    if ws_comment is not None:
        with open(ws_comment, 'w') as f:
            f.write("Compilation-performance test failed")

    # First _get_ all the sources (with as little a gap between
    # fetches as possible to minimize race possibility)
    for instance in instances:
        workspace = get_workspace_for_instance(instance, args)

        if args.setup_workspaces_for_pr:
            setup_workspace(instance, workspace, args)
        validate_workspace(instance, workspace, args)

    # Then _build and test_ both workspaces, which can take
    # longer since we're not concerned about racing on github
    # anymore.
    for instance in instances:
        workspace = get_workspace_for_instance(instance, args)

        if not args.skip_build:
            build_swift_toolchain(workspace, args)

        if not args.skip_runner:
            execute_runner(instance, workspace, configs, args)

    regressions = analyze_results(configs, args)
    return regressions


def get_sandbox_profile_flags():
    sandbox_flags = []
    if 'WORKSPACE' not in os.environ:
        raise RuntimeError("Missing $WORKSPACE when looking for sandbox")
    if 'JOB_NAME' not in os.environ:
        raise RuntimeError("Missing $JOB_NAME when looking for sandbox")
    workspace = os.environ['WORKSPACE']
    job_name = os.environ['JOB_NAME']
    sscss = os.path.join(workspace, "../../workspace-private",
                         job_name, "swift-source-compat-suite-sandbox")
    if not os.path.exists(sscss):
        raise RuntimeError("Missing sandbox dir %s" % sscss)

    if platform.system() == 'Darwin':
        sandbox_flags += [
            '--sandbox-profile-xcodebuild',
            os.path.join(sscss, 'sandbox_xcodebuild.sb'),
            '--sandbox-profile-package',
            os.path.join(sscss, 'sandbox_package.sb')
        ]
    elif platform.system() == 'Linux':
        sandbox_flags += [
            '--sandbox-profile-package',
            os.path.join(sscss, 'sandbox_package_linux.profile')
        ]
    else:
        raise common.UnsupportedPlatform
    return sandbox_flags


def get_swiftc_path(instance, workspace, args):
    if instance == OLD_INSTANCE and args.old_swiftc is not None:
        return args.old_swiftc
    elif instance == NEW_INSTANCE and args.new_swiftc is not None:
        return args.new_swiftc
    elif platform.system() == 'Darwin':
        swiftc_path = os.path.join(workspace, 'build/compat_macos/install/toolchain/usr/bin/swiftc')
    elif platform.system() == 'Linux':
        swiftc_path = os.path.join(workspace, 'build/compat_linux/install/usr/bin/swiftc')
    else:
        raise common.UnsupportedPlatform
    return swiftc_path


def get_preset_name(args):
    platform_name = None
    if platform.system() == 'Darwin':
        platform_name = 'macos'
    elif platform.system() == 'Linux':
        platform_name = 'linux'
    else:
        raise common.UnsupportedPlatform
    assert(platform_name is not None)

    build_type = ''
    if args.debug:
        build_type = build_type + 'D'
    else:
        build_type = build_type + 'R'

    if args.assertions:
        build_type = build_type + 'A'

    preset = 'source_compat_suite_%s_%s'
    return preset % (platform_name, build_type)


def build_swift_toolchain(workspace, args):
    build_script_preset = get_preset_name(args)

    build_script_args_common = [
        '--preset=%s' % build_script_preset
    ]

    if args.distcc:
        build_script_args_common += ['--distcc']
    if args.cmake_c_launcher:
        build_script_args_common += ['--cmake-c-launcher={}'.format(args.cmake_c_launcher)]
    if args.cmake_cxx_launcher:
        build_script_args_common += ['--cmake-cxx-launcher={}'.format(args.cmake_cxx_launcher)]

    build_command = [os.path.join(workspace, 'swift/utils/build-script')]
    build_command += build_script_args_common

    if platform.system() == 'Darwin':
        build_command += [
            'install_destdir={}/build/compat_macos/install'.format(workspace),
            'install_prefix=/toolchain/usr',
            'install_symroot={}/build/compat_macos/symroot'.format(workspace),
            'installable_package={}/build/compat_macos/root.tar.gz'.format(workspace),
            'symbols_package={}/build/compat_macos/root-symbols.tar.gz'.format(workspace),
        ]
    elif platform.system() == 'Linux':
        build_command += [
            'install_destdir={}/build/compat_linux/install'.format(workspace),
            'install_prefix=/usr',
            'installable_package={}/build/compat_linux/root.tar.gz'.format(workspace),
        ]
    else:
        raise common.UnsupportedPlatform
    common.check_execute(build_command, timeout=9999999)


def get_projects(suite):
    if suite == 'full':
        return 'projects.json'
    elif suite == 'smoketest':
        return 'projects-cperf-smoketest.json'
    else:
        raise ValueError("Unknown suite: " + suite)


def get_configs(suite):
    if suite == 'full':
        return ['debug-batch', 'release']
    elif suite == 'smoketest':
        return ['debug', 'release']
    else:
        raise ValueError("Unknown suite: " + suite)


def get_stats_dir(instance, variant):
    return os.path.join(os.getcwd(),
                        "-".join(["stats", instance, variant]))


def get_actual_config_and_flags(config, stats):
    flags = ("-stats-output-dir  " + stats)
    # Handle pseudo-configs
    if config == 'wmo-onone':
        flags += ' -wmo -Onone '
        config = 'release'
    elif config == 'debug-opt':
        flags += ' -O '
        config = 'debug'
    elif config == 'debug-batch':
        flags += ' -enable-batch-mode '
        config = 'debug'
    return (config, flags)


def get_variant(config, args):
    return '-'.join([args.suite, args.swift_branch, config])


def execute_runner(instance, workspace, configs, args):
    projects = get_projects(args.suite)
    swiftc_path = get_swiftc_path(instance, workspace, args)
    for config in configs:
        variant = get_variant(config, args)
        stats = get_stats_dir(instance, variant)
        if os.path.exists(stats):
            shutil.rmtree(stats)
        os.makedirs(stats)
        (config, flags) = get_actual_config_and_flags(config, stats)
        actions = 'action.startswith("BuildSwiftPackage")'
        if platform.system() == 'Darwin':
            # Darwin can handle xcodebuild as well
            actions = 'action.startswith("Build")'
        runner_command = [
            './runner.py',
            '--swiftc', swiftc_path,
            '--projects', projects,
            '--build-config', config,
            '--include-actions', actions,
            '--swift-branch', args.swift_branch,
            '--add-swift-flags', flags,
        ]
        if args.sandbox:
            runner_command += get_sandbox_profile_flags()
        if args.verbose:
            runner_command += ["--verbose"]
        for i in range(args.repetitions):
            try:
                common.check_execute(runner_command, timeout=9999999)
            except common.ExecuteCommandFailure:
                pass


def get_table_name(reference, subset, variant):
    return os.path.join(
        os.getcwd(),
        ('-'.join([reference, subset, variant]) + '.md'))


def get_config_desc(config):
    return config.capitalize()


def get_reference_desc(reference, config):
    return ("PR vs. %s (%s)" %
            (reference, config))


def get_table_desc(refqual, subset, config):
    if refqual is not None:
        return ("PR vs. %s, changed %s (%s)" %
                (refqual, subset, config))
    else:
        return ("%s %s" % (config, subset))


def make_internal_link(desc, run_id):
    link = re.sub('[^ a-zA-Z0-9-]', '', desc)
    link = link.strip().lower()
    link = re.sub(' +', '-', link)
    link += '-' + run_id
    return ("[%s](#%s)" % (desc, link))


def make_internal_anchor(desc, run_id):
    link = re.sub('[^ a-zA-Z0-9-]', '', desc)
    link = link.strip().lower()
    link = re.sub(' +', '-', link)
    link += '-' + run_id
    return ('\n<a name="%s">\n' % link)


def get_baseline_name(variant):
    return os.path.join(os.getcwd(),
                        'cperf-baselines', variant + '.csv')


def all_driver_stats_files(configs, args):
    for config in configs:
        variant = get_variant(config, args)
        for instance in [OLD_INSTANCE, NEW_INSTANCE]:
            for root, dirs, files in os.walk(get_stats_dir(instance, variant)):
                for f in files:
                    m = re.match(r'stats-\d+-swift-driver-([^-]+).*\.json', f)
                    if m:
                        try:
                            with open(os.path.join(root, f)) as fp:
                                j = json.load(fp)
                            yield (config, instance, m.group(1), j)
                        except ValueError:
                            # Something sufficiently bad happened that
                            # we can't read the json: build a fake one.
                            yield (config, instance, m.group(1), {})


def find_module_statuses(configs, args):
    passed = set()
    failed = set()
    k = 'Driver.NumProcessFailures'
    module_sets = {}
    for config in configs:
        module_sets[config] = {}
        module_sets[config][OLD_INSTANCE] = set()
        module_sets[config][NEW_INSTANCE] = set()

    for (config, instance, module, j) in all_driver_stats_files(configs, args):
        # Treat a module as failed if there was >0 process failures,
        # or we can't tell, _or_ it's in one instance but not in the other.
        module_sets[config][instance].add(module)
        if j.get(k, 1) == 0:
            passed.add(module)
        else:
            failed.add(module)

    mismatched_modules = None
    for config in configs:
        for instance in [OLD_INSTANCE, NEW_INSTANCE]:
            mods = module_sets[config][instance]
            if mismatched_modules is None:
                mismatched_modules = mods
            else:
                mismatched_modules.symmetric_difference_update(mods)
    failed.update(mismatched_modules)
    passed.difference_update(mismatched_modules)
    return (passed, failed)


def analyze_results(configs, args):
    # Make a separate random run_id for each comment, to uniquify links.
    run_id = hex(random.randint(1, 2**63))[2:]
    old_ws = get_workspace_for_instance(OLD_INSTANCE, args)
    process_stats = os.path.join(old_ws, 'swift/utils/process-stats-dir.py')
    (passed_modules, failed_modules) = find_module_statuses(configs, args)
    references = ('head', 'baseline') if args.show_baselines else ('head',)
    subsets = ('brief', 'detailed')
    delta_usec_thresh = str(100 * 1000)
    delta_pct_thresh = str(1)
    common_args = [process_stats,
                   '--markdown', "--github-emoji"]
    if args.group_by_module:
        common_args += ['--group-by-module']

    for m in (passed_modules - failed_modules):
        common_args += ['--select-module', m]

    briefpattern = r'driver.*wall|NumLLVMBytesOutput|NumInstructions'
    returncodes = []
    for config in configs:
        variant = get_variant(config, args)
        sd_old = get_stats_dir(OLD_INSTANCE, variant)
        sd_new = get_stats_dir(NEW_INSTANCE, variant)
        baseline = get_baseline_name(variant)
        select_stats_from_baseline_args = (
            ["--select-stats-from-csv-baseline", baseline]
            if os.path.exists(baseline)
            else [])
        returncodes.append(
            common.execute(
                common_args +
                ['--output', get_table_name('head', 'brief', variant),
                 '--divide-by', str(args.repetitions),
                 '--select-stat', briefpattern,
                 '--merge-timers',
                 '--delta-pct-thresh', delta_pct_thresh,
                 '--delta-usec-thresh', delta_usec_thresh,
                 '--compare-stats-dirs', sd_old, sd_new]))
        returncodes.append(
            common.execute(
                common_args +
                select_stats_from_baseline_args +
                ['--output', get_table_name('head', 'detailed', variant),
                 '--exclude-timers',
                 '--divide-by', str(args.repetitions),
                 '--close-regressions',
                 '--delta-pct-thresh', delta_pct_thresh,
                 '--compare-stats-dirs', sd_old, sd_new]))
        if args.show_baselines and os.path.exists(baseline):
            returncodes.append(
                common.execute(
                    common_args +
                    ['--output',
                     get_table_name('baseline', 'brief', variant),
                     '--select-stat', briefpattern,
                     '--merge-timers',
                     '--divide-by', str(args.repetitions),
                     '--delta-pct-thresh', delta_pct_thresh,
                     '--delta-usec-thresh', delta_usec_thresh,
                     '--compare-to-csv-baseline', baseline, sd_new]))
            returncodes.append(
                common.execute(
                    common_args +
                    ['--output',
                     get_table_name('baseline', 'detailed', variant),
                     '--exclude-timers',
                     '--divide-by', str(args.repetitions),
                     '--close-regressions',
                     '--delta-pct-thresh', delta_pct_thresh,
                     '--compare-to-csv-baseline', baseline, sd_new]))

    out = args.output
    regressions = any(x != 0 for x in returncodes)
    out.write("# Summary for %s %s\n\n" % (args.swift_branch, args.suite))
    if len(failed_modules) != 0:
        out.write("**Unexpected test results, excluded stats for %s**\n\n" %
                  ", ".join(failed_modules))
    if regressions:
        out.write("**Regressions found (see below)**\n\n")
    else:
        out.write("No regressions above thresholds\n\n")

    # Write summary TOC
    for config in configs:
        variant = get_variant(config, args)
        cdesc = get_config_desc(config)
        out.write('- %s\n' % make_internal_link(cdesc, run_id))
        for reference in references:
            rindent = ''
            refqual = None
            if len(references) > 1:
                rindent = '    '
                refqual = reference
                rdesc = get_reference_desc(reference, config)
                out.write('    - %s\n' % make_internal_link(rdesc, run_id))
            for subset in subsets:
                tdesc = get_table_desc(refqual, subset, config)
                out.write((rindent + '    - %s\n') %
                          make_internal_link(tdesc, run_id))

    out.write('\n\n')

    for config in configs:
        variant = get_variant(config, args)
        cdesc = get_config_desc(config)
        out.write(make_internal_anchor(cdesc, run_id))
        out.write("\n# %s\n" % cdesc)
        for reference in references:
            rindent = ''
            refqual = None
            if len(references) > 1:
                rindent = '#'
                refqual = reference
                rdesc = get_reference_desc(reference, config)
                out.write(make_internal_anchor(rdesc, run_id))
                out.write("\n## %s\n" % rdesc)
            for subset in subsets:
                tdesc = get_table_desc(refqual, subset, config)
                out.write(make_internal_anchor(tdesc, run_id))
                out.write("\n" + rindent + "## " + tdesc + "\n")
                table = get_table_name(reference, subset, variant)
                if os.path.exists(table):
                    if os.stat(table).st_size == 0:
                        out.write("\nNone\n")
                    else:
                        with open(table) as t:
                            out.write(t.read())
                    os.unlink(table)
                else:
                    out.write("\nNo analysis available\n")

    if args.show_baselines:
        out.write('\n\n<details>\n')
        for config in configs:
            variant = get_variant(config, args)
            baseline = get_baseline_name(variant)
            if os.path.exists(baseline):
                out.write("Last baseline commit on %s" %
                          os.path.basename(baseline))
                out.write("\n\n<pre>\n")
                out.write(common.check_execute_output([
                    'git', 'log', '--pretty=medium', '-1', baseline
                ]))
                out.write("\n</pre>\n")
            else:
                out.write("No baseline file %s found" %
                          os.path.basename(baseline))
        out.write('\n</details>\n\n')

    return regressions


def get_default_output(basename):
    if 'WORKSPACE' in os.environ:
        workspace = os.environ['WORKSPACE']
        return os.path.join(workspace, basename)
    else:
        return basename


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('swift_branch')
    parser.add_argument('--sandbox', action='store_true')
    parser.add_argument('--url', type=str,
                        default="git@github.com:apple/swift.git")
    parser.add_argument('--suite',
                        metavar='(smoketest|full)',
                        help='cperf suite to run',
                        default='smoketest')
    parser.add_argument("--verbose",
                        action='store_true',
                        default=False)
    parser.add_argument("--quiet",
                        action='store_false',
                        dest='verbose')
    parser.add_argument('--skip-build',
                        action='store_true')
    parser.add_argument('--skip-runner',
                        action='store_true')
    parser.add_argument('--output',
                        type=argparse.FileType('w'),
                        default=get_default_output('comment.md'))
    parser.add_argument('--clean',
                        action='store_true')
    parser.add_argument('--setup-workspaces-for-pr',
                        type=int, default=None)
    parser.add_argument('--old-workspace',
                        type=str, default=None)
    parser.add_argument('--new-workspace',
                        type=str, default=None)
    parser.add_argument('--old-swiftc',
                        type=str, default=None)
    parser.add_argument('--new-swiftc',
                        type=str, default=None)
    parser.add_argument('--repetitions',
                        type=int, default=2)
    parser.add_argument('--show-baselines',
                        type=bool, default=False)
    parser.add_argument('--group-by-module',
                        type=bool, default=False)
    parser.add_argument('--self-test',
                        action='store_true')
    parser.add_argument('--cmake-c-launcher',
                        metavar="PATH",
                        help='the absolute path to set CMAKE_C_COMPILER_LAUNCHER for build script')
    parser.add_argument('--cmake-cxx-launcher',
                        metavar='PATH',
                        help='the absolute path to set CMAKE_CXX_COMPILER_LAUNCHER for build script')
    parser.add_argument("--distcc",
                        help='Build Swift pass distcc as cmake-*-launcher to build-script',
                        action='store_true')
    parser.add_argument("--debug",
                        help='Build Swift in debug mode',
                        action='store_true')
    parser.add_argument("--assertions",
                        help='Build Swift with assertions enabled',
                        action='store_true')
    return parser.parse_args()


if __name__ == '__main__':
    sys.exit(main())
