#!/usr/bin/env python3
#
# Copyright 2018 Ettus Research, a National Instruments Company
#
# SPDX-License-Identifier: LGPL-3.0-or-later
#

import argparse
import os
import sys
import subprocess
import logging
import re
import time
import datetime

from junit_xml import TestSuite, TestCase, to_xml_report_file
from queue import Queue
from threading import Thread

# -------------------------------------------------------
# Utilities
# -------------------------------------------------------

SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
BASE_DIR = os.path.split(os.path.split(SCRIPT_DIR)[0])[0]

_LOG = logging.getLogger(os.path.basename(__file__))
_LOG.setLevel(logging.INFO)
_STDOUT = logging.StreamHandler()
_LOG.addHandler(_STDOUT)
_FORMATTER = logging.Formatter("[%(name)s] - %(levelname)s - %(message)s")
_STDOUT.setFormatter(_FORMATTER)

RETCODE_SUCCESS = 0
RETCODE_PARSE_ERR = -1
RETCODE_EXEC_ERR = -2
RETCODE_COMPILE_ERR = -3
RETCODE_UNKNOWN_ERR = -4


def retcode_to_str(code):
    """Convert internal status code to string"""
    code = int(code)
    if code > RETCODE_SUCCESS:
        return "AppError({code})".format(code=code)
    else:
        return {
            RETCODE_SUCCESS: "OK",
            RETCODE_PARSE_ERR: "ParseError",
            RETCODE_EXEC_ERR: "ExecError",
            RETCODE_COMPILE_ERR: "CompileError",
            RETCODE_UNKNOWN_ERR: "UnknownError",
        }[code]


def result_to_string(result):
    """Convert a result dictionary to a string summarizing the results"""
    if "module" in result:
        retval = "Expected={:02d}, Run={:02d}, Passed={:02d}, Elapsed={:s}, Returncode={:s}".format(
            result["tc_expected"],
            result["tc_run"],
            result["tc_passed"],
            result["wall_time"],
            retcode_to_str(result["retcode"]),
        )
    else:
        retval = "Returncode={:s}".format(retcode_to_str(result["retcode"]))
    return retval


def log_with_header(what, minlen=0, ch="#"):
    """Print with a header around the text"""
    padlen = max(int((minlen - len(what)) / 2), 1)
    toprint = (" " * padlen) + what + (" " * padlen)
    _LOG.info(ch * len(toprint))
    _LOG.info(toprint)
    _LOG.info(ch * len(toprint))


# -------------------------------------------------------
# Simulation Functions
# -------------------------------------------------------


def read_excludes_file(excludes_fname):
    if excludes_fname:
        return [l.strip() for l in open(excludes_fname) if (l.strip() and "#" not in l)]
    else:
        return []


def find_sims_on_fs(basedir, excludes):
    """Find all testbenches in the specific basedir
    Testbenches are defined as directories with a
    Makefile that includes viv_sim_preamble.mak
    """
    sims = {}
    for root, _, files in os.walk(basedir):
        name = os.path.relpath(root, basedir)
        if "Makefile" in files:
            with open(os.path.join(root, "Makefile"), "r") as mfile:
                for l in mfile.readlines():
                    if re.match(".*include.*viv_sim_preamble.mak.*", l) is not None:
                        if name not in excludes:
                            sims.update({name: root})
                            break
    return sims


def gather_target_sims(basedir, targets, excludes):
    """Parse the specified targets and gather simulations to run
    Remove duplicates and sort alphabetically
    """
    fs_sims = find_sims_on_fs(basedir, excludes)
    if not isinstance(targets, list):
        targets = [targets]
    sim_names = set()
    for target in targets:
        for name in sorted(fs_sims):
            if re.match(target, name) is not None:
                sim_names.add(name)
    target_sims = []
    for name in sorted(sim_names):
        target_sims.append((name, fs_sims[name]))
    return target_sims


def parse_output(simout):
    # Gather results (basic metrics)
    results = {"retcode": RETCODE_SUCCESS, "stdout": simout, "passed": False}
    # Look for the following in the log:
    # - A start timestamp (indicates that Vivado started)
    # - The testbench infrastructure start header (indicates that the TB started)
    # - A stop timestamp (indicates that the TB stopped)
    tb_started = False
    compile_started = False
    results["start_time"] = "<unknown>"
    results["wall_time"] = "<unknown>"
    for line in simout.split(b"\n"):
        tsm = re.match(rb"TESTBENCH STARTED: (.+)", line)
        if tsm is not None:
            tb_started = True
        csm = re.match(rb"source .*viv_sim_project.tcl", line)
        if csm is not None:
            compile_started = True
        # xsim log
        vsm = re.match(rb"# Start of session at: (.+)", line)
        if vsm is not None:
            results["start_time"] = str(vsm.group(1), "ascii")
        tfm = re.match(rb"launch_simulation:.*; elapsed = (.+) \..*", line)
        if tfm is not None:
            results["wall_time"] = str(tfm.group(1), "ascii")
        # modelsim log
        vsm = re.match(rb"# Start time: (.+)", line)
        if vsm is not None:
            results["start_time"] = str(vsm.group(1), "ascii")
        tfm = re.match(rb"# End time:.*, Elapsed time: (.+)", line)
        if tfm is not None:
            results["wall_time"] = str(tfm.group(1), "ascii")
    # Parse testbench results
    # We have two possible formats to parse because we have two simulation
    # test executors.
    tb_match_fmt0 = [
        b".*TESTBENCH FINISHED: (.+)\n",
        b"(?:# )? - Time elapsed:   (.+) ns.*\n",
        b"(?:# )? - Tests Expected: (.+)\n",
        b"(?:# )? - Tests Run:      (.+)\n",
        b"(?:# )? - Tests Passed:   (.+)\n",
        b"(?:# )?Result: (PASSED|FAILED).*",
    ]
    m_fmt0 = re.match(b"".join(tb_match_fmt0), simout, re.DOTALL)
    tb_match_fmt1 = [
        b".*TESTBENCH FINISHED: (.*)\n",
        b"(?:# )? - Time elapsed:  (.+) ns.*\n",
        b"(?:# )? - Tests Run:     (.+)\n",
        b"(?:# )? - Tests Passed:  (.+)\n",
        b"(?:# )? - Tests Failed:  (.+)\n",
        b"(?:# )?Result: (PASSED|FAILED).*",
    ]
    m_fmt1 = re.match(b"".join(tb_match_fmt1), simout, re.DOTALL)

    # Remove escape characters (colors) from Vivado output
    ansi_escape = re.compile(r"(?:\x1B[\(@-_]|[\x80-\x9F])[0-?]*[ -/]*[@-~]")
    plain_simout = ansi_escape.sub("", simout.decode("utf-8"))

    # Check for $error() and $fatal() output, which may be missed by the
    # testbench or may occur in a subsequent instance, after a pass.
    tb_match_error = [
        "\n",
        "(Error|Fatal): .*\n.?",
        "Time: .+\n",
    ]
    m_error = re.search("".join(tb_match_error), plain_simout)

    # Figure out the returncode
    retcode = RETCODE_UNKNOWN_ERR
    if m_fmt0 is not None or m_fmt1 is not None:
        retcode = RETCODE_SUCCESS
        if m_fmt0 is not None:
            results["passed"] = m_fmt0.group(6) == b"PASSED" and m_error is None
            results["module"] = m_fmt0.group(1)
            results["sim_time_ns"] = int(m_fmt0.group(2))
            results["tc_expected"] = int(m_fmt0.group(3))
            results["tc_run"] = int(m_fmt0.group(4))
            results["tc_passed"] = int(m_fmt0.group(5))
        else:
            results["passed"] = m_fmt1.group(6) == b"PASSED" and m_error is None
            results["module"] = m_fmt1.group(1)
            results["sim_time_ns"] = int(m_fmt1.group(2))
            results["tc_expected"] = int(m_fmt1.group(3))
            results["tc_run"] = int(m_fmt1.group(3))
            results["tc_passed"] = int(m_fmt1.group(4))
    elif tb_started:
        retcode = RETCODE_PARSE_ERR
    elif compile_started:
        retcode = RETCODE_COMPILE_ERR
    else:
        retcode = RETCODE_EXEC_ERR
    results["retcode"] = retcode
    return results


def run_sim(path, simulator, basedir, setupenv):
    """Run the simulation at the specified path
    The simulator can be specified as the target
    A environment script can be run optionally
    """
    try:
        # Optionally run an environment setup script
        if setupenv is None:
            setupenv = ""
            # Check if environment was setup
            if "VIVADO_PATH" not in os.environ:
                return {
                    "retcode": RETCODE_EXEC_ERR,
                    "passed": False,
                    "stdout": bytes("Simulation environment was not initialized\n", "utf-8"),
                }
        else:
            setupenv = ". " + os.path.realpath(setupenv) + ";"
        # Run the simulation
        return parse_output(
            subprocess.check_output(
                'cd {workingdir}; /bin/bash -c "{setupenv} make ip 2>&1; make {simulator} 2>&1"'.format(
                    workingdir=path, setupenv=setupenv, simulator=simulator
                ),
                shell=True,
            )
        )
    except subprocess.CalledProcessError as e:
        return {"retcode": int(abs(e.returncode)), "passed": False, "stdout": e.output}
    except Exception as e:
        _LOG.error("Target " + path + " failed to run:\n" + str(e))
        return {"retcode": RETCODE_EXEC_ERR, "passed": False, "stdout": bytes(str(e), "utf-8")}
    except:
        _LOG.error("Target " + path + " failed to run")
        return {
            "retcode": RETCODE_UNKNOWN_ERR,
            "passed": False,
            "stdout": bytes("Unknown Exception", "utf-8"),
        }


def run_sim_queue(run_queue, out_queue, simulator, basedir, setupenv):
    """Thread worker for a simulation runner
    Pull a job from the run queue, run the sim, then place
    output in out_queue
    """
    while not run_queue.empty():
        (name, path) = run_queue.get()
        try:
            _LOG.info("Starting: %s", name)
            result = run_sim(path, simulator, basedir, setupenv)
            out_queue.put((name, result))
            _LOG.info(
                "FINISHED: %s (%s, %s)",
                name,
                retcode_to_str(result["retcode"]),
                "PASS" if result["passed"] else "FAIL!",
            )
        except KeyboardInterrupt:
            _LOG.warning("Target " + name + " received SIGINT. Aborting...")
            out_queue.put(
                (
                    name,
                    {
                        "retcode": RETCODE_EXEC_ERR,
                        "passed": False,
                        "stdout": bytes("Aborted by user", "utf-8"),
                    },
                )
            )
        except Exception as e:
            _LOG.error("Target " + name + " failed to run:\n" + str(e))
            out_queue.put(
                (
                    name,
                    {
                        "retcode": RETCODE_UNKNOWN_ERR,
                        "passed": False,
                        "stdout": bytes(str(e), "utf-8"),
                    },
                )
            )
        finally:
            run_queue.task_done()


# -------------------------------------------------------
# Script Actions
# -------------------------------------------------------


def do_list(args):
    """List all simulations that can be run"""
    excludes = read_excludes_file(args.excludes)
    for (name, path) in gather_target_sims(args.basedir, args.target, excludes):
        print(name)
    return 0


def do_run(args):
    """Build a simulation queue based on the specified
    args and process it
    """
    run_queue = Queue(maxsize=0)
    out_queue = Queue(maxsize=0)
    _LOG.info("Queueing the following targets to simulate:")
    excludes = read_excludes_file(args.excludes)
    name_maxlen = 0
    for (name, path) in gather_target_sims(args.basedir, args.target, excludes):
        run_queue.put((name, path))
        name_maxlen = max(name_maxlen, len(name))
        _LOG.info("* " + name)
    # Spawn tasks to run builds
    num_sims = run_queue.qsize()
    num_jobs = min(num_sims, int(args.jobs))
    _LOG.info("Started " + str(num_jobs) + " job(s) to process queue...")
    results = {}
    for i in range(num_jobs):
        worker = Thread(
            target=run_sim_queue,
            args=(run_queue, out_queue, args.simulator, args.basedir, args.setupenv),
        )
        worker.setDaemon(False)
        worker.start()
    # Wait for build queue to become empty
    start = datetime.datetime.now()
    try:
        sim_count = -1
        while out_queue.qsize() < num_sims:
            tdiff = str(datetime.datetime.now() - start).split(".", 2)[0]
            if args.logged:
                # Print number of TBs completed and elapsed time whenever a
                # simulation completes.
                if sim_count != out_queue.qsize():
                    print(
                        ">>> [%s] (%d/%d simulations completed) <<<"
                        % (tdiff, out_queue.qsize(), num_sims)
                    )
                    sim_count = out_queue.qsize()
            else:
                # Print elapsed time and number of TBs completed once per
                # second, overwriting the same line each time.
                print(
                    "\r>>> [%s] (%d/%d simulations completed) <<<"
                    % (tdiff, out_queue.qsize(), num_sims),
                    end="\r",
                    flush=True,
                )
            time.sleep(1.0)
        sys.stdout.write("\n")
    except (KeyboardInterrupt):
        _LOG.warning("Received SIGINT. Aborting... (waiting for pending jobs to finish)")
        # Flush run queue
        while not run_queue.empty():
            (name, path) = run_queue.get()
        raise SystemExit(1)

    results = {}
    result_all = 0
    while not out_queue.empty():
        (name, result) = out_queue.get()
        results[name] = result
        line = "#" * 70
        sys.stdout.buffer.write(bytes("%s\n Begin TB Log: %s\n%s\n" % (line, name, line), "utf-8"))
        sys.stdout.buffer.write(result["stdout"])
        sys.stdout.buffer.write(bytes("%s\n End TB Log: %s\n%s\n" % (line, name, line), "utf-8"))
        if not result["passed"]:
            result_all += 1
    sys.stdout.write("\n\n\n")
    sys.stdout.flush()
    time.sleep(1.0)

    hdr_len = name_maxlen + 62  # 62 is the report line length
    log_with_header("RESULTS", hdr_len)
    for name in sorted(results):
        result = results[name]
        _LOG.info(
            "* %s : %s (%s)",
            name.ljust(name_maxlen),
            ("Passed" if result["passed"] else "FAILED"),
            result_to_string(result),
        )
    _LOG.info("=" * hdr_len)
    _LOG.info(
        "SUMMARY: %d out of %d tests passed. Time elapsed was %s"
        % (num_sims - result_all, num_sims, str(datetime.datetime.now() - start).split(".", 2)[0])
    )
    _LOG.info("#" * hdr_len)
    if args.report:
        do_report(args, results)
    return result_all


def do_cleanup(args):
    """Run make cleanall for all simulations"""
    setupenv = args.setupenv
    if setupenv is None:
        setupenv = ""
        # Check if environment was setup
        if "VIVADO_PATH" not in os.environ:
            raise RuntimeError("Simulation environment was not initialized")
    else:
        setupenv = ". " + os.path.realpath(setupenv) + ";"
    excludes = read_excludes_file(args.excludes)
    for (name, path) in gather_target_sims(args.basedir, args.target, excludes):
        _LOG.info("Cleaning up %s", name)
        os.chdir(os.path.join(args.basedir, path))
        subprocess.Popen("{setupenv} make cleanall".format(setupenv=setupenv), shell=True).wait()
    return 0


def do_report_csv(args, results):
    """Generate report file (.csv)"""
    keys = [
        "module",
        "status",
        "retcode",
        "start_time",
        "wall_time",
        "sim_time_ns",
        "tc_expected",
        "tc_run",
        "tc_passed",
    ]
    with open(args.report, "w") as repfile:
        repfile.write((",".join([x.upper() for x in keys])) + "\n")
        for name in sorted(results):
            r = results[name]
            if r["retcode"] != RETCODE_SUCCESS:
                results["retcode"] = retcode_to_str(r["retcode"])
                results["status"] = "ERROR"
                results["start_time"] = r["start_time"]
            else:
                results = r
                results["module"] = name
                results["status"] = "PASSED" if r["passed"] else "FAILED"
                results["retcode"] = retcode_to_str(r["retcode"])
            repfile.write((",".join([str(results[x]) for x in keys])) + "\n")
    _LOG.info("Testbench report (CSV format) written to " + args.report)
    return 0


def do_report_junit(args, results):
    """Generate JUnit report file (.xml)"""
    test_cases = []
    for name in sorted(results):
        result = results[name]
        if "wall_time" in result:
            h, m, s = map(int, result["wall_time"].split(":"))
            elapsed_sec = h * 3600 + m * 60 + s
        else:
            elapsed_sec = None
        test_case = TestCase(name=name, elapsed_sec=elapsed_sec)
        if result["retcode"] != RETCODE_SUCCESS:
            test_case.add_error_info(
                message=retcode_to_str(result["retcode"]), output=result["stdout"].decode()
            )
        test_cases.append(test_case)
    test_suite = TestSuite(name=f"{args.simulator} testbenches", test_cases=test_cases)
    with open(args.report, "w") as repfile:
        to_xml_report_file(repfile, [test_suite], prettyprint=True)
    _LOG.info("Testbench report (JUnit format) written to " + args.report)
    return 0


def do_report(args, results):
    """Generate report file (either .csv or .xml)"""

    def _load_results(args):
        results = {}
        excludes = read_excludes_file(args.excludes)
        for (name, path) in gather_target_sims(args.basedir, args.target, excludes):
            logpath = os.path.join(path, args.simulator + ".log")
            if os.path.isfile(logpath):
                with open(logpath, "rb") as logfile:
                    result = parse_output(logfile.read())
                    results[name] = result
        return results

    suffix = os.path.splitext(args.report)[1]
    if results is None:
        results = _load_results(args)
    if suffix == ".csv":
        return do_report_csv(args, results)
    elif suffix == ".xml":
        return do_report_junit(args, results)
    else:
        raise ValueError(f"Unsupported report suffix: {suffix}")


# Parse command line options
def get_options():
    parser = argparse.ArgumentParser(description="Batch testbench execution script")
    parser.add_argument(
        "-d",
        "--basedir",
        default=BASE_DIR,
        help="Base directory in which to search for testbenches",
    )
    parser.add_argument(
        "-s",
        "--simulator",
        choices=["xsim", "vsim", "modelsim"],
        default="xsim",
        help="Simulator name",
    )
    parser.add_argument(
        "-e",
        "--setupenv",
        default=None,
        help="Optional environment setup script to run for each TB",
    )
    parser.add_argument(
        "-r", "--report", default="testbench_report.csv", help="Name of the output report file"
    )
    parser.add_argument(
        "-x",
        "--excludes",
        default=None,
        help="Name of the excludes file. It contains all targets to exclude.",
    )
    parser.add_argument("-j", "--jobs", default=1, help="Number of parallel simulation jobs to run")
    parser.add_argument(
        "-l",
        "--logged",
        action="store_true",
        default=False,
        help="Output is logged, so don't show per-second timer",
    )
    parser.add_argument(
        "action", choices=["run", "cleanup", "list", "report"], default="list", help="What to do?"
    )
    parser.add_argument(
        "target", nargs="*", default=".*", help="Space separated simulation target regexes"
    )
    return parser.parse_args()


def main():
    args = get_options()
    actions = {"list": do_list, "run": do_run, "cleanup": do_cleanup, "report": do_report}
    return actions[args.action](args)


if __name__ == "__main__":
    exit(main())