import enum
import sys
import os

from amaranth.build import Platform
from amaranth import *
from amaranth.sim import Simulator

class ClockState(enum.Enum):
    LOW     = 0
    FALLING = 1
    HIGH    = 2
    RISING  = 3

class ExiClock(Elaboratable):

    def __init__(self):
        # Ports
        self.exiClk = Signal()
        self.exiClkState = Signal(2)

        # State
        self.prevExiClkValid = Signal()
        self.prevExiClk = Signal()
        self.prevExiClkState = Signal(2)

    
    def elaborate(self, platform: Platform):
        m = Module()

        with m.If(self.prevExiClkValid):
            with m.If(self.prevExiClkState == ClockState.FALLING):
                m.d.comb += self.exiClkState.eq(ClockState.LOW)
            with m.Elif(self.prevExiClkState == ClockState.RISING):
                m.d.comb += self.exiClkState.eq(ClockState.HIGH)
            with m.Else():
                with m.If(self.prevExiClk ^ self.exiClk):
                    m.d.comb += self.exiClkState.eq(Cat(1, self.exiClk))
                with m.Else():
                    m.d.comb += self.exiClkState.eq(Cat(0, self.exiClk))

            m.d.sync += self.prevExiClkState.eq(self.exiClkState)
            
        m.d.sync += self.prevExiClk.eq(self.exiClk)
        m.d.sync += self.prevExiClkValid.eq(1)
        
        return m

#####
# TestBench
#####
class TestBench:
    def __init__(self):
        pass

    def FlipExiClock(self, dut):
        yield dut.exiClk.eq(~dut.exiClk)

    def clockTest(self):
        dut = self.dut

        yield dut.exiClk.eq(0)
        yield
        yield from self.FlipExiClock(dut)
        yield
        assert (yield dut.exiClkState) == ClockState.RISING.value
        yield
        assert (yield dut.exiClkState) == ClockState.HIGH.value
        yield
        yield from self.FlipExiClock(dut)
        yield
        assert (yield dut.exiClkState) == ClockState.FALLING.value
        yield
        assert (yield dut.exiClkState) == ClockState.LOW.value
        yield


    def simulate(self):
        self.dut = ExiClock()

        sim = Simulator(self.dut)
        sim.add_clock(1e-6) # 1 MHz
        sim.add_sync_process(self.clockTest)

        with sim.write_vcd(os.path.dirname(os.path.abspath(__file__)) + "/ExiClock.vcd"):
            sim.run()

#####
# Main portion
#####

def main():
    if(len(sys.argv) == 2):

        if sys.argv[1] == "s":
            bench = TestBench()
            bench.simulate()
        
        if sys.argv[1] == "v":
            mod = ExiClock()
            with open(os.path.dirname(os.path.abspath(__file__)) + "/ExiClock.v", "w") as f:
                f.write(verilog.convert(mod, ports=[mod.exiClk, mod.exiClkState]))
    
    else:
        bench = TestBench()
        bench.simulate()

if __name__ == "__main__":
    main()