#!/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")