import sys
from amaranth import *
from amaranth.back import verilog, cxxrtl
from amaranth.cli import main
from amaranth.sim import Simulator, Settle, Delay

BASENAME = "shift_reg"

class ShiftReg(Elaboratable):
    def __init__(self, width):
        self.load_val = Signal(width, reset=0, reset_less=True)
        self.load = Signal()
        self.reg = Signal(width)
        self.en = Signal()
        self.right_left = Signal()

        self.ports = [self.load_val, self.en, self.right_left, self.reg]

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

        with m.If(self.load):
            m.d.sync += self.reg.eq(self.load_val)
        with m.Else():
            with m.If(self.en):
                with m.If(self.right_left):
                    m.d.sync += self.reg.eq(self.reg << 1)
                with m.Else():
                    m.d.sync += self.reg.eq(self.reg >> 1)

        return m


def step():
    yield
    yield Settle()

def test_shift_reg():
    dut = ShiftReg(8)

    def proc1():
        val = 0xAB

        yield dut.load_val.eq(val)
        yield dut.en.eq(0)
        yield dut.load.eq(1)
        yield
        yield Settle()
        yield dut.load.eq(0)
        yield dut.en.eq(1)   

        for _ in range(9):
            reg_val = yield dut.reg
            assert reg_val == val, f"Incorrect shift ---EXPECTED: {hex(val)}   ---GOT: {hex(reg_val)}"
            val = val >> 1
            yield
            yield Settle()

        val = 0xBD
        yield dut.load_val.eq(val)
        yield dut.load.eq(1)
        yield dut.right_left.eq(1)
        yield
        yield Settle()
        yield dut.load.eq(0)

        for _ in range(9):
            reg_val = yield dut.reg
            assert reg_val == val, f"Incorrect shift ---EXPECTED: {hex(val)}   ---GOT: {hex(reg_val)}"
            val = (val << 1) & 0xff
            yield
            yield Settle()



    sim = Simulator(dut)
    sim.add_clock(1e-6)
    sim.add_sync_process(proc1)
    
    with sim.write_vcd(BASENAME + '.vcd'):
        sim.run()


if __name__ == '__main__':

    if sys.argv[1] == "sim":
        test_shift_reg()
        exit()

    m = ShiftReg(8)

    if sys.argv[1] == "v":
        out = verilog.convert(m, ports=m.ports)
        with open(BASENAME + '.v','w') as f:
            f.write(out)
    
    elif sys.argv[1] == "cc":
        out = cxxrtl.convert(m, ports=m.ports)
        with open(BASENAME + '.cc','w') as f:
            f.write(out)