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, invert=False):
        self.gpio_out = gpio_out
        self.count = count
        self.invert = 1 if invert else 0

        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)) + 1,
                      data_width=bus_gr, type="mux") as self.csr:
            self.csr.r += csr.Register("switch", "rw", width=count)
            with self.csr.r.switch as reg:
                reg.f += [
                    csr.Field("output_{}".format(i), reset_value=self.invert)
                    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] ^ self.invert)
            for i in range(self.count)
        ]

        return m