"""
Set of utilities to build a simple test suite.
"""
from amaranth import *
from amaranth.sim import *

from typing import Generator
import unittest
import os

from contextlib import nullcontext

class BaseTestClass(unittest.TestCase):
    """
    Base test class that provides a run_test helper function to do all the nice things.
    """

    def _run_test(self, test: Generator, name: str):
        try:
            sim = Simulator(self.harness)
        except NameError:
            raise NotImplementedError(f"Must define a self.harness module for TestCase {self.__class__.__name__}!")

        sim.add_clock(100e-9)
        sim.add_sync_process(test)
        sim.reset()

        # Pretty hacky way to pass this info in but does it look like I care?
        if os.environ.get("TEST_SAVE_VCD"):
           ctx = sim.write_vcd(f"vcd_out/{name}.vcd")
        else:
            ctx = nullcontext()

        with ctx:
            sim.run()

        del sim


    ######### Random Utilities ########
    def _write_csr(self, bus, index, data):
        yield bus.addr.eq(index)
        yield bus.w_stb.eq(1)
        yield bus.w_data.eq(data)
        yield Tick()
        yield bus.w_stb.eq(0)
        yield Tick()


    def _wait_for_signal(self, signal, polarity=False, require_edge=True, timeout=1000):
        ready_for_edge = not require_edge # If we don't require edge, we can just ignore

        while True:
            timeout -= 1
            if timeout == 0:
                self.fail(f"_wait_for_signal({signal}, {polarity}, {require_edge}, {timeout}, timed out!")

            read = yield signal
            if read == polarity:
                if ready_for_edge:
                    break
            else:
                ready_for_edge = True

            yield Tick()


def provide_testcase_name(fn):
    """Decorator that provides a function with access to its own class and name."""
    def wrapper(self):
        fn(self, f"{self.__class__.__name__}.{fn.__name__}")

    return wrapper