outdir = "submission"
include_dirs = ["./modules", "./tests"]
src_dirs = ["tests"]
icarus_verilog_args = ["-g2012"]

import argparse
import subprocess
from pathlib import Path

parser = argparse.ArgumentParser(prog="ISV",
                        description="A wrapper around iverilog and vvp to provide useful project utilities")
parser.add_argument("cmd")

src_files: list[Path] = []
for src_dir in src_dirs:
    src_dir = Path(src_dir)
    if not (src_dir.is_file() or src_dir.is_dir()):
        print(f"\033[38;5;9m\"{str(src_dir)}\" Is not a file or directory\033[m")
        continue
    src_files += [file.absolute().relative_to(Path.cwd()) for file in src_dir.rglob("*") if file.is_file()]


def subcmd_submit(args):
    for src_file in src_files:
        if src_file.name == "tests.sv" or src_file.name == "tests.v":
            continue

        Path.mkdir(Path(outdir), exist_ok=True)
        prog_args = [
            "iverilog", *icarus_verilog_args, "-E", "-o", str(Path(outdir) / Path(src_file.name))
        ]
        for include_dir in include_dirs:
            prog_args += [f"-I{include_dir}"]
        prog_args.append(str(src_file))
        print(prog_args)
        subprocess.run(prog_args)

def subcmd_test_submit(args):
    test_files: list[Path] = [file for file in Path(outdir).rglob("*") if file.is_file()]
    failed_tests = 0
    failed_compiles = 0
    for test_file in test_files:
        prev_failed_compiles = failed_compiles
        failed_compiles += int(subprocess.run(["iverilog", *icarus_verilog_args, "-o", "testbench.vvp", str(test_file)]).returncode != 0)
        if prev_failed_compiles == failed_compiles:
            failed_tests += int(subprocess.run(["vvp", "testbench.vvp"]).returncode != 0)
    if failed_tests > 0 or failed_compiles > 0:
        print(f"\033[38;5;9mFailed {failed_tests} tests with {failed_compiles} tests not compiling for submission\033[m")
    else: print(f"\033[38;5;10mAll submission tests passing\033[m")

def subcmd_test(args):
    prog_args = [
        "iverilog", *icarus_verilog_args, "-o", "testbench.vvp"
    ]
    for include_dir in include_dirs:
        prog_args += [f"-I{include_dir}"]
    prog_args.append("tests/tests.sv")
    if subprocess.run(prog_args).returncode == 0:
        assert subprocess.run(["vvp", "testbench.vvp"]).returncode == 0, "Failed to run testbench program"
    else: print(f"\033[38;5;9mFailed to compile testbench program\033[m")

subcommands = {
    "submit": subcmd_submit,
    "test_submit": subcmd_test_submit,
    "test": subcmd_test
}


if __name__ == "__main__":
    args = parser.parse_args()
    assert args.cmd in subcommands, f"Invalid sub-command \"{args.cmd}\""
    subcommands[args.cmd](args)