import os
import sys

import argparse

from cocotb.runner import get_runner

from rtl_manifest import rtl_manifest

import yaml


def fpga_sim_main():
    # 1: Parse Arguments.

    parser = argparse.ArgumentParser(
        prog="sim",
        description="Tool to simulate"
    )

    parser.add_argument("yaml")
    parser.add_argument("test_name", nargs="?")
    parser.add_argument("-j", "--jobs", default=1)

    args = parser.parse_args()

    print(args.test_name)

    # print(args.yaml)

    with open(args.yaml) as cfg_file:
        cfg = yaml.safe_load(cfg_file)

    # 2: Figure out which tests to run

    base_path = os.path.split(os.path.abspath(args.yaml))[0]

    tests = []

    defines = {}

    def parse_cfg(_cfg, _base_path=None):
        if "yaml" in _cfg:
            for sub_yaml in _cfg["yaml"]:
                print(sub_yaml)
                with open(f"{_base_path}/{sub_yaml}") as _cfg_file:
                    __cfg = yaml.safe_load(_cfg_file)
                    print(f"Base Path: {_base_path}/{os.path.split(sub_yaml)[0]}")
                    parse_cfg(__cfg, f"{_base_path}/{os.path.split(sub_yaml)[0]}")

        if "tests" in _cfg:
            for test in _cfg["tests"]:
                print(test)
                if "waves" not in test:
                    waves = False
                else:
                    waves = test["waves"]
                
                tests.append({
                    "base_path": _base_path,
                    "name": test["name"],
                    "toplevel": test["toplevel"],
                    "modules": test["modules"],
                    "waves": waves,
                    "sources": test["sources"]
                })

        if "defines" in _cfg:
            cfg_defines = _cfg["defines"]
            for define in cfg_defines:
                print(f"{define}: {cfg_defines[define]}, {os.path.expandvars(cfg_defines[define])}")
                defines[define] = os.path.expandvars(cfg_defines[define])

    parse_cfg(cfg, base_path)

    # 4: Run those tests
    sim = os.getenv("SIM", "verilator")
    runner = get_runner(sim)

    jobs = args.jobs
    print(jobs)
    os.environ["MAKEFLAGS"] = f"-j{jobs}"

    tests_to_run = []

    if args.test_name is not None:
        for test in tests:
            if test["name"] == args.test_name:
                tests_to_run.append(test)
    else:
        tests_to_run = tests

    # Turn this into a multiprocessing pool
    for test in tests_to_run:

        sources, incdirs = rtl_manifest.parse(f"{test['base_path']}/{test['sources']}")

        runner.build(
            verilog_sources=sources,
            includes=incdirs,
            hdl_toplevel=test["toplevel"],
            build_dir=f"{test['base_path']}/sim_build",
            waves=test["waves"],
            defines=defines
        )

        result_xml = f"../sim_build/{test['name']}_results.xml".replace(" ", "_")

        sys.path.append(test["base_path"])

        runner.test(hdl_toplevel=test["toplevel"], test_module=test["modules"], waves=test["waves"], results_xml=result_xml)