#!/usr/bin/env python
# utils/coverage/coverage-generate-data - Generate, parse test run profdata
#
# This source file is part of the Swift.org open source project
#
# Copyright (c) 2014 - 2017 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

import argparse
import logging
import multiprocessing
import os
import pipes
import platform
import subprocess
import sys
import timeit
from multiprocessing import Pool

NUM_CORES = multiprocessing.cpu_count()

logging_format = '%(asctime)s %(levelname)s %(message)s'
logging.basicConfig(level=logging.DEBUG,
                    format=logging_format,
                    filename='/tmp/%s.log' % os.path.basename(__file__),
                    filemode='w')
console = logging.StreamHandler()
console.setLevel(logging.INFO)
formatter = logging.Formatter(logging_format)
console.setFormatter(formatter)
logging.getLogger().addHandler(console)

global_build_subdir = ''


def quote_shell_cmd(cmd):
    """Return `cmd` as a properly quoted shell string"""
    return ' '.join([pipes.quote(a) for a in cmd])


def call(cmd, verbose=True, show_cmd=True):
    """Call `cmd` and optionally log debug info"""
    formatted_cmd = quote_shell_cmd(cmd) if isinstance(cmd, list) else cmd
    if show_cmd:
        logging.info('$ ' + formatted_cmd)
    start_time = timeit.default_timer()
    process = subprocess.Popen(
        cmd,
        shell=(not isinstance(cmd, list)),
        bufsize=1,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT
    )
    for line in iter(process.stdout.readline, b''):
        if verbose:
            logging.info('STDOUT: ' + line.rstrip())
    end_time = timeit.default_timer()
    logging.debug('END $ ' + formatted_cmd)
    logging.debug('Return code: %s', process.returncode)
    logging.debug('Elapsed time: %s', end_time - start_time)
    return process.returncode


def check_output(cmd, verbose=True, show_cmd=True):
    """Return output of calling `cmd` and optionally log debug info"""
    output = []
    formatted_cmd = quote_shell_cmd(cmd) if isinstance(cmd, list) else cmd
    if show_cmd:
        logging.info('$ ' + formatted_cmd)
    start_time = timeit.default_timer()
    process = subprocess.Popen(
        cmd,
        shell=(not isinstance(cmd, list)),
        bufsize=1,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT
    )
    for line in iter(process.stdout.readline, b''):
        if verbose:
            logging.info('STDOUT: ' + line.rstrip())
        output.append(line)
    end_time = timeit.default_timer()
    logging.debug('Return code: %s', process.returncode)
    logging.debug('Elapsed time: %s', end_time - start_time)
    return (process.returncode, ''.join(output))


def xcrun_find(cmd):
    """Return path of `cmd` using xcrun -f"""
    return check_output(['xcrun', '-f', cmd])[1].strip()


llvm_cov = xcrun_find('llvm-cov')
llvm_profdata = xcrun_find('llvm-profdata')


def dump_coverage_data(merged_file):
    """Dump coverage data of file at path `merged_file` using llvm-cov"""
    try:
        swift = os.path.join(global_build_subdir,
                             'swift-macosx-{}/bin/swift'.format(platform.machine()))
        coverage_log = os.path.join(os.path.dirname(merged_file),
                                    'coverage.log')
        testname = os.path.basename(os.path.dirname(merged_file))
        logging.info('Searching for covered files: %s', testname)
        (returncode, output) = check_output(
            [llvm_cov, 'report', '-instr-profile=%s' % merged_file, swift],
            verbose=False,
            show_cmd=False
        )
        output = [line.split()[0]
                  for line in output.split()
                  if '0.00' not in line and '/swift' in line]
        with open(coverage_log, 'w') as f:
            logging.info('Dumping coverage data: %s', testname)
            (returncode2, dumped) = check_output(
                quote_shell_cmd(
                    [llvm_cov, 'show', '-line-coverage-gt=0',
                     '-instr-profile=%s' % merged_file, swift] +
                    output
                ),
                verbose=False,
                show_cmd=False
            )
            f.write(dumped)
    except Exception as e:
        logging.debug(str(e))


def find_folders(root_path, suffix):
    """Return a list of folder paths ending in `suffix` rooted at
    `root_path`"""
    found_folders = []
    for root, folders, files in os.walk(root_path):
        for folder in folders:
            if folder.endswith(suffix):
                folderpath = os.path.join(root, folder)
                logging.debug('Found %s', folderpath)
                found_folders.append(folderpath)
    logging.info('Found %s "%s" folders', len(found_folders), suffix)
    return found_folders


def find_files(root_path, suffix):
    """Return a list of file paths ending in `suffix` rooted at
    `root_path`"""
    found_files = []
    for root, folders, files in os.walk(root_path):
        for f in files:
            if f.endswith(suffix):
                fpath = os.path.join(root, f)
                logging.debug('Found %s', fpath)
                found_files.append(fpath)
    logging.info('Found %s "%s" files', len(found_files), suffix)
    return found_files


def merge_profdir(profdir_path):
    """Merge swift-*.profraw files contained in `profdir_path` into
    merged.profraw"""
    logging.info('Merging %s', profdir_path)
    if not os.path.exists(os.path.join(profdir_path, 'merged.profraw')):
        call('set -x; '
             'cd %s; '
             '%s merge -output merged.profraw swift-*.profraw && '
             'rm swift-*.profraw' % (profdir_path, llvm_profdata))


def demangle_coverage_data(coverage_log_path):
    """Demangle coverage dump at `coverage_log_path` using c++filt"""
    logging.info('Demangling %s', coverage_log_path)
    cppfilt = '/usr/bin/c++filt'
    demangled_log_path = coverage_log_path + '.demangled'
    returncode = 1
    with open(coverage_log_path) as cf, open(demangled_log_path, 'w') as df:
        process = subprocess.Popen(
            [cppfilt, '-n'],
            stdin=subprocess.PIPE,
            stdout=df,
            stderr=subprocess.PIPE
        )
        for line in cf:
            process.stdin.write(line)
        process.stdin.close()
        returncode = process.wait()
    return returncode


def main():
    global global_build_subdir

    parser = argparse.ArgumentParser(
        description='Generate, parse test run profdata')
    parser.add_argument('swift_dir', metavar='swift-dir')
    parser.add_argument('--log',
                        help='the level of information to log (default: info)',
                        metavar='LEVEL',
                        default='info',
                        choices=['info', 'debug', 'warning', 'error',
                                 'critical'])
    args = parser.parse_args()

    console.setLevel(level=args.log.upper())
    logging.debug(args)

    swift_dir = os.path.realpath(os.path.abspath(args.swift_dir))
    build_dir = os.path.realpath(os.path.join(os.path.dirname(swift_dir),
                                              'build'))
    build_subdir = os.path.join(build_dir, 'buildbot_incremental_coverage')

    global_build_subdir = build_subdir

    build_script_cmd = [
        os.path.join(swift_dir, 'utils/build-script'),
        '--preset=buildbot_incremental,tools=RDA,stdlib=RDA,coverage',
    ]

    call(build_script_cmd)

    assert global_build_subdir

    pool = Pool(NUM_CORES)

    logging.info('Starting merge on %s', build_dir)
    folders = find_folders(build_dir, '.profdir')
    pool.map_async(merge_profdir, folders).get(999999)

    logging.info('Starting coverage data dump...')
    merged_profraw_files = find_files(build_dir, 'merged.profraw')
    pool.map_async(dump_coverage_data, merged_profraw_files).get(999999)

    logging.info('Starting coverage data dump demangling...')
    coverage_log_files = find_files(build_dir, 'coverage.log')
    pool.map_async(demangle_coverage_data, coverage_log_files).get(999999)
    return 0


if __name__ == '__main__':
    try:
        sys.exit(main())
    except Exception as e:
        logging.debug(str(e))
        sys.exit(1)
