import sys
import os
from amaranth import *
from amaranth.back import verilog
from amaranth.build.plat import Platform
from amaranth.sim import Simulator
from amaranth_boards.icebreaker import *

from ExiClock import ExiClock
from ShiftRegister import ShiftRegister

class ReBba(Elaboratable):
    def __init__(self):
        pass

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

        led1 = platform.request("led_g", 1)
        led2 = platform.request("led_g", 4)
        led3 = platform.request("led_g", 2)
        led4 = platform.request("led_g", 3)
        led5 = platform.request("led_r", 1)

        btn1 = platform.request("button", 1)
        btn2 = platform.request("button", 2)
        btn3 = platform.request("button", 3)

        exiClk = ExiClock()
        m.submodules += exiClk
        sr = ShiftRegister(5)
        m.submodules += sr

        m.d.comb += led1.eq(sr.data[0])
        m.d.comb += led2.eq(sr.data[1])
        m.d.comb += led3.eq(sr.data[2])
        m.d.comb += led4.eq(sr.data[3])
        m.d.comb += led5.eq(sr.data[4])

        m.d.comb += exiClk.exiClk.eq(btn3)

        m.d.comb += sr.rst.eq(btn2)
        m.d.comb += sr.inb.eq(btn1)
        m.d.comb += sr.nen.eq(Const(0))
        m.d.comb += sr.exiClkState.eq(exiClk.exiClkState)

        return m

class TestBench:
    def __init__(self):
        pass

    def simulate(self):
        pass

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

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

        if sys.argv[1] == "p":
            mod = ReBba()
            platform = ICEBreakerPlatform()
            platform.add_resources(platform.break_off_pmod)
            platform.build(mod, do_program=True)
    
    else:
        bench = TestBench()
        bench.simulate()


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

if __name__ == "__main__":
    main()