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


__all__ = ["TimerCore"]


class TimerCore(Elaboratable):
    def __init__(self, *, width, bus_data_width,
                 clk_domain="sync", bus_granularity=None):
        self.width = width

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

        with csr.Bank(name="timer", 
                      addr_width=log2_int(-(-(bus_dw+3*max(width, bus_dw))//bus_gr), need_pow2=False), 
                      data_width=bus_gr) as self.csr:
            self.csr.r += [
                csr.Register(
                    "control", "rw", width=bus_dw,
                    fields=[csr.Field("enable", startbit=0),
                            csr.Field("update_enable", startbit=4)]
                ),
                csr.Register(
                    "load", "rw", width=max(width, bus_dw), 
                    fields=[csr.Field("value", width=self.width)]
                ),
                csr.Register(
                    "reload", "rw", width=max(width, bus_dw), 
                    fields=[csr.Field("value", width=self.width)]
                ),
                csr.Register(
                    "current", "r", width=max(width, bus_dw),
                    fields=[csr.Field("value", width=self.width)]
                )
            ]
        self.wb2csr = csr.WishboneCSRBridge(self.csr.dec.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

        # Value right now
        value = Signal(self.width)
        # When Timer is enabled, start counting
        with m.If(self.csr.r.control.f.enable.s):
            with m.If(value == 0):
                m.d.sync += value.eq(self.csr.r.reload.s)
            with m.Else():
                m.d.sync += value.eq(value - 1)
        # Otherwise, load Timer value to .value
        with m.Else():
            m.d.sync += value.eq(self.csr.r.load.s)

        # If Timer allows updating value in memory, update
        m.d.comb += self.csr.r.current.set_val.eq(value)
        with m.If(self.csr.r.control.f.update_enable.s):
            m.d.sync += self.csr.r.current.set_stb.eq(1)
        with m.Else():
            m.d.sync += self.csr.r.current.set_stb.eq(0)

        return m