#! /usr/bin/env python3

from multiprocessing import Pool
import multiprocessing
import argparse
import tempfile
import logging
import os
import subprocess


def run_reproducer(path):
    proc = subprocess.Popen([LLDB, '--replay', path],
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE)
    reason = None
    try:
        outs, errs = proc.communicate(timeout=TIMEOUT)
        success = proc.returncode == 0
        result = 'PASSED' if success else 'FAILED'
        if not success:
            outs = outs.decode()
            errs = errs.decode()
            # Do some pattern matching to find out the cause of the failure.
            if 'Encountered unexpected packet during replay' in errs:
                reason = 'Unexpected packet'
            elif 'Assertion failed' in errs:
                reason = 'Assertion failed'
            elif 'UNREACHABLE' in errs:
                reason = 'Unreachable executed'
            elif 'Segmentation fault' in errs:
                reason = 'Segmentation fault'
            elif 'Illegal instruction' in errs:
                reason = 'Illegal instruction'
            else:
                reason = f'Exit code {proc.returncode}'
    except subprocess.TimeoutExpired:
        proc.kill()
        success = False
        outs, errs = proc.communicate()
        result = 'TIMEOUT'

    if not FAILURE_ONLY or not success:
        reason_str = f' ({reason})' if reason else ''
        print(f'{result}: {path}{reason_str}')
        if VERBOSE:
            if outs:
                print(outs)
            if errs:
                print(errs)


def find_reproducers(path):
    for root, dirs, files in os.walk(path):
        for dir in dirs:
            _, extension = os.path.splitext(dir)
            if dir.startswith('Test') and extension == '.py':
                yield os.path.join(root, dir)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='LLDB API Test Replay Driver. '
        'Replay one or more reproducers in parallel using the specified LLDB driver. '
        'The script will look for reproducers generated by the API lit test suite. '
        'To generate the reproducers, pass --param \'lldb-run-with-repro=capture\' to lit.'
    )
    parser.add_argument(
        '-j',
        '--threads',
        type=int,
        default=multiprocessing.cpu_count(),
        help='Number of threads. The number of CPU threads if not specified.')
    parser.add_argument(
        '-t',
        '--timeout',
        type=int,
        default=60,
        help='Replay timeout in seconds. 60 seconds if not specified.')
    parser.add_argument(
        '-p',
        '--path',
        type=str,
        default=os.getcwd(),
        help=
        'Path to the directory containing the reproducers. The current working directory if not specified.'
    )
    parser.add_argument('-l',
                        '--lldb',
                        type=str,
                        required=True,
                        help='Path to the LLDB command line driver')
    parser.add_argument('-v',
                        '--verbose',
                        help='Print replay output.',
                        action='store_true')
    parser.add_argument('--failure-only',
                        help='Only log failures.',
                        action='store_true')
    args = parser.parse_args()

    global LLDB
    global TIMEOUT
    global VERBOSE
    global FAILURE_ONLY
    LLDB = args.lldb
    TIMEOUT = args.timeout
    VERBOSE = args.verbose
    FAILURE_ONLY = args.failure_only

    print(
        f'Replaying reproducers in {args.path} with {args.threads} threads and a {args.timeout} seconds timeout'
    )

    try:
        pool = Pool(args.threads)
        pool.map(run_reproducer, find_reproducers(args.path))
    except KeyboardInterrupt:
        print('Interrupted')