#!/usr/bin/env python3

import os
import subprocess
import sys

args = sys.argv[1:]

generate_vcd = False
if "--vcd" in args:
	generate_vcd = True
	del args[args.index("--vcd")]

if len(args) > 0:
	testlist = args
	# This happens a lot when autocomplete is used:
	for i, n in enumerate(testlist):
		if n.endswith(".c"):
			testlist[i] = n[:-2]
else:
	testlist = []
	for path in os.listdir():
		if os.path.isfile(path) and path.endswith(".c"):
			testlist.append(path[:-2])

testlist = sorted(testlist)

tb_build_ret = subprocess.run(
	["make", "-C", "../tb_cxxrtl", "tb"],
	timeout=300
)
if tb_build_ret.returncode != 0:
	sys.exit("Failed.")

all_passed = True

for test in testlist:
	sys.stdout.write(f"{test:<30}")

	test_build_ret = subprocess.run(
		["make", f"APP={test}", f"tmp/{test}.bin"],
		stdout=subprocess.DEVNULL
	)
	if test_build_ret.returncode != 0:
		print("\033[33m[MK ERR]\033[39m")
		all_passed = False
		continue

	cmdline = ["../tb_cxxrtl/tb", "--bin", f"tmp/{test}.bin", "--cycles", "1000000"]
	if generate_vcd:
		cmdline += ["--vcd", f"tmp/{test}.vcd"]

	try:
		test_run_ret = subprocess.run(
			cmdline,
			stdout = subprocess.PIPE,
			stderr = subprocess.PIPE,
			timeout=10
		)
		with open(f"tmp/{test}.log", "wb") as f:
			f.write(test_run_ret.stdout)
	except subprocess.TimeoutExpired:
		print("\033[31m[TIMOUT]\033[39m")
		all_passed = False
		continue

	# Testbench itself should always exit successfully.
	if test_run_ret.returncode != 0:
		print("Negative return code from testbench!")
		all_passed = False
		continue

	# Pass if the program under test has zero exit code AND its output matches
	# the expected output (if there is an expected_output file)

	output_lines = test_run_ret.stdout.decode("utf-8").strip().splitlines()
	returncode = -1
	if len(output_lines) >= 2:
		exit_line = output_lines[-2]
		if exit_line.startswith("CPU requested halt"):
			try:
				returncode = int(exit_line.split(" ")[-1])
			except:
				pass
	if returncode != 0:
		print("\033[31m[BADRET]\033[39m")
		all_passed = False
		continue

	test_src = open(f"{test}.c").read()
	if "/*EXPECTED-OUTPUT" in test_src:
		good_output = True
		try:
			expected_start = test_src.find("/*EXPECTED-OUTPUT")
			expected_end = test_src.find("*/", expected_start)
			expected_lines = test_src[expected_start:expected_end + 1].splitlines()[1:-1]
			while expected_lines[0].strip() == "":
				del expected_lines[0]
			while expected_lines[-1].strip() == "":
				del expected_lines[-1]

			# Allow single-line comments within the expected output, in case some of
			# the output needs explanation inline in the test source. If the line is
			# empty after stripping comments, still don't remove the line.
			for i, l in enumerate(expected_lines):
				if "//" in l:
					expected_lines[i] = l.split("//")[0].rstrip()

			# Drop last two lines, which should just be tb output (checked in BADRET)
			output_lines = output_lines[:-2]
			while output_lines[0].strip() == "":
				del output_lines[0]
			while output_lines[-1].strip() == "":
				del output_lines[-1]

			if expected_lines != output_lines:
				good_output = False
		except:
			good_output = False
		if not good_output:
			print("\033[31m[BADOUT]\033[39m")
			all_passed = False
			continue

	print("\033[32m[PASSED]\033[39m")

sys.exit(not all_passed)