#/usr/bin/env python

"""
This script is used to provide a wrapper to difftest, so that they can be easily
deployed on ci environment.
"""

import os
import os.path as osp
import sys
from multiprocessing import Pool, Process
from functools import partial

def find_all_test_images(path):
    tests = []
    for root, dirs, files in os.walk(path):
        for file in files:
            # Get file extensions and select files with .bin
            ext = osp.splitext(file)[1]
            if ext == ".bin":
                tests.append(osp.join(root, file))
    return tests

def run_test(test_image, ref, ref_prefix, dut, dut_prefix):
    diffu = "diffu"
    args = [
        "--images-path", "/",
        "-m", test_image,
        "--ref", ref,
        "--ref-prefix", ref_prefix,
        "--dut", dut,
        "--dut-prefix", dut_prefix,
        "> logs/" + osp.basename(test_image) + ".log",
        "2>&1"
    ]

    status = os.system(diffu + " " + " ".join(args))
    exitcode = os.waitstatus_to_exitcode(status)

    image_shortname = osp.basename(test_image)
    print(f"{ 'FAILED' if exitcode else 'PASSED' } {image_shortname}")
    if exitcode:
      print(f"cmd: {diffu + ' ' + ' '.join(args)}")
      print(f"exitcode: {exitcode}")
    sys.exit(exitcode)

def print_statistics(results):
    pass

def main():
    DIFFU_IMAGES_PATH = os.environ["DIFFU_IMAGES_PATH"]
    print(DIFFU_IMAGES_PATH)
    assert(osp.isdir(DIFFU_IMAGES_PATH))
    os.makedirs("logs", exist_ok = True)
    # Run tests in a multiprocess pool
    tests = find_all_test_images(DIFFU_IMAGES_PATH)
    ref, ref_prefix, dut, dut_prefix = sys.argv[1:]
    ref_shortname = osp.basename(ref)
    dut_shortname = osp.basename(dut)
    print(f"[{ref_shortname}, {dut_shortname}]")
    procs = []
    for test in tests:
      image_shortname = osp.basename(test)
      p = Process(target=run_test, args=(test, ref, ref_prefix, dut, dut_prefix), name=image_shortname, daemon=True)
      procs.append(p)
      p.start()

    timeout = 0
    for p in procs:
      p.join(5)
      if p.exitcode is None:
        print(f"{ 'TIMEOUT' } {p.name}")
        p.terminate()
        p.join()
        timeout += 1

    not_success = sum((1 for p in procs if p.exitcode != 0))
    failed = not_success - timeout

    print("==========")
    print(f"TOTAL {len(procs)}\tFAILED: {failed}\tTIMEOUT: {timeout}")

    return not_success

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