#!/usr/bin/env python3
# ===----------------------------------------------------------------------===##
#
# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#
# ===----------------------------------------------------------------------===##

import argparse
import json
import logging
import os
import pathlib
import platform
import subprocess
import sys
import tempfile


def directory_path(string):
    if os.path.isdir(string):
        return pathlib.Path(string)
    else:
        raise NotADirectoryError(string)

def gather_machine_information(args):
    """
    Gather the machine information to upload to LNT as part of the submission.
    """
    info = {}
    if platform.system() == 'Darwin':
        profiler_info = json.loads(subprocess.check_output(['system_profiler', 'SPHardwareDataType', 'SPSoftwareDataType', '-json']).decode())
        info['hardware'] = profiler_info['SPHardwareDataType'][0]['chip_type']
        info['os'] = profiler_info['SPSoftwareDataType'][0]['os_version']
        info['sdk'] = subprocess.check_output(['xcrun', '--show-sdk-version']).decode().strip()

    info['compiler'] = subprocess.check_output([args.compiler, '--version']).decode().strip().splitlines()[0]
    info['test_suite_commit'] = subprocess.check_output(['git', '-C', args.git_repo, 'rev-parse', args.test_suite_commit]).decode().strip()

    if args.spec_dir is not None:
        with open(args.spec_dir / 'version.txt', 'r') as f:
            info['spec'] = f.read().strip()

    return info

def gather_run_information(args):
    """
    Gather the run information to upload to LNT as part of the submission.
    """
    info = {}
    # TODO: For now, we don't include commit information. This happens to be a long string and
    #       production instances of LNT dont't accept those as part of submissions for the time being.
    # info['commit_info'] = subprocess.check_output(['git', '-C', args.git_repo, 'show', args.benchmark_commit, '--no-patch']).decode()
    info['git_sha'] = subprocess.check_output(['git', '-C', args.git_repo, 'rev-parse', args.benchmark_commit]).decode().strip()
    return info

def dict_to_params(d):
    """
    Return a list of 'key=value' strings from a dictionary.
    """
    res = []
    for (k, v) in d.items():
        res.append(f'{k}={v}')
    return res

def main(argv):
    parser = argparse.ArgumentParser(
        prog='run-benchmarks',
        description='Benchmark libc++ at the given commit and produce a LNT JSON report.')
    parser.add_argument('--benchmark-commit', type=str, required=True,
        help='The SHA representing the version of the library to benchmark.')
    parser.add_argument('--test-suite-commit', type=str, required=True,
        help='The SHA representing the version of the test suite to use for benchmarking.')
    parser.add_argument('--compiler', type=str, required=True,
        help='Path to the compiler to use for testing.')
    parser.add_argument('--machine', type=str, required=True,
        help='The name of the machine for reporting LNT results.')
    parser.add_argument('--output', type=pathlib.Path, required=True,
        help='Path where the resulting LNT JSON report is written. The file is overwritten if it already exists.')
    parser.add_argument('--filter', type=str, required=False,
        help="Optional test filter to pass to lit when running the benchmarks. This allows "
             "running only a subset of the benchmarks.")
    parser.add_argument('--spec-dir', type=pathlib.Path, required=False,
        help='Optional path to a SPEC installation to use for benchmarking.')
    parser.add_argument('--git-repo', type=directory_path, default=os.getcwd(),
        help='Optional path to the Git repository to use. By default, the current working directory is used.')
    parser.add_argument('--dry-run', action='store_true',
        help='Do not actually perform any action. Use with -vv to see what would be executed.')
    parser.add_argument('-v', '--verbose', action='count', default=0,
        help='Verbosity level: passing the option multiple times increases the level.')
    args = parser.parse_args(argv)

    if args.verbose == 0:
        logging.basicConfig(level=logging.INFO)
    elif args.verbose >= 1:
        logging.basicConfig(level=logging.DEBUG)

    def run(command, *posargs, enforce_success=True, **kwargs):
        command = [str(c) for c in command]
        logging.debug(f'$ {" ".join(command)}')
        if args.dry_run:
            return

        # If we are running with verbose, don't capture any output: just let the subprocess
        # print anything to stdout and stderr. Otherwise, capture stderr and stdout so we can
        # diagnose when something goes wrong.
        try:
            if not args.verbose:
                if 'stdout' not in kwargs:
                    kwargs['stdout'] = subprocess.PIPE
                if 'stderr' not in kwargs:
                    kwargs['stderr'] = subprocess.PIPE
            subprocess.run(command, check=True, *posargs, **kwargs)
        except subprocess.CalledProcessError as e:
            if e.stdout:
                sys.stdout.write(e.stdout.decode())
            if e.stderr:
                sys.stderr.write(e.stderr.decode())
            # If we enforce success, don't swallow the exception
            if enforce_success:
                raise

    with tempfile.TemporaryDirectory() as build_dir:
        build_dir = pathlib.Path(build_dir)

        logging.info('Installing LNT')
        run(['python3', '-m', 'venv', build_dir / '.venv'])
        run([build_dir / '.venv/bin/pip', 'install', 'llvm-lnt'])

        logging.info(f'Building libc++ at commit {args.benchmark_commit}')
        build_cmd = [args.git_repo / 'libcxx/utils/build-at-commit',
                        '--git-repo', args.git_repo,
                        '--install-dir', build_dir / 'install',
                        '--commit', args.benchmark_commit,
                        '--', '-DCMAKE_BUILD_TYPE=RelWithDebInfo', f'-DCMAKE_CXX_COMPILER={args.compiler}']
        run(build_cmd, enforce_success=False) # if the build fails, carry on: we'll fail later and submit empty LNT results

        logging.info(f'Running benchmarks from {args.test_suite_commit} against libc++ {args.benchmark_commit}')
        cmd = [args.git_repo / 'libcxx/utils/test-at-commit',
                            '--git-repo', args.git_repo,
                            '--build-dir', build_dir / 'bench',
                            '--test-suite-commit', args.test_suite_commit,
                            '--libcxx-installation', build_dir / 'install',
                            '--',
                            '-j1', '--time-tests', '--test-output=failed',
                            '--param', f'compiler={args.compiler}',
                            '--param', 'optimization=speed',
                            '--param', 'std=c++26',
                            build_dir / 'bench/libcxx/test/benchmarks']
        if args.spec_dir is not None:
            cmd += ['--param', f'spec_dir={args.spec_dir}']
        if args.filter is not None:
            cmd += ['--filter', args.filter]
        run(cmd, enforce_success=False) # some benchmarks may fail to build/run at some commits, and that's okay
        with open(build_dir / 'benchmarks.lnt', 'w') as f:
            run([args.git_repo / 'libcxx/utils/consolidate-benchmarks', build_dir / 'bench'], stdout=f)

        logging.info('Creating JSON report for LNT')
        order = len(subprocess.check_output(['git', '-C', args.git_repo, 'rev-list', args.benchmark_commit]).splitlines())
        importreport = [build_dir / '.venv/bin/lnt', 'importreport', '--order', str(order), '--machine', args.machine]
        for arg in dict_to_params(gather_run_information(args)):
            importreport += ['--run-info', arg]
        for arg in dict_to_params(gather_machine_information(args)):
            importreport += ['--machine-info', arg]
        output = args.output.resolve()
        importreport += [build_dir / 'benchmarks.lnt', output]
        run(importreport)
        logging.info(f'Report written to {output}')


if __name__ == '__main__':
    main(sys.argv[1:])
