from nmigen import *
from nmigen_soc import csr
from nmigen.utils import *


__all__ = ["GPIOOutput"]


class GPIOOutput(Elaboratable):
    def __init__(self, gpio_out, *, bus_data_width, count=8, bus_granularity=None):
        self.gpio_out = gpio_out
        self.count = count

        bus_dw = bus_data_width
        if bus_granularity is None:
            bus_granularity = data_width
        bus_gr = bus_granularity

        with csr.Bank(name="gpio",
                      addr_width=max(1, log2_int(-(-count//bus_gr), need_pow2=False)),
                      data_width=bus_gr, type="mux") as self.csr:
            self.csr.r += [
                csr.Register(
                    "switch", "rw", width=count,
                    fields=[
                        csr.Field("output_{}".format(i)) for i in range(count)
                    ]
                )
            ]
        self.wb2csr = csr.WishboneCSRBridge(self.csr.mux.bus, data_width=bus_data_width)
        self.csr_bus = self.wb2csr.wb_bus

    def elaborate(self, platform):
        m = Module()
        
        m.submodules += self.csr, self.wb2csr

        m.d.comb += [
            self.gpio_out[i].eq(self.csr.r.switch.s[i])
            for i in range(self.count)
        ]

        return m