#!/usr/bin/env python3

from amaranth import *
from amaranth.sim import *
from amaranth_boards import colorlight_i9
from amaranth_soc.wishbone import Interface, Arbiter, Decoder
from amaranth_soc.memory import MemoryMap

from typing import List

from minerva.core import Minerva

class Blinky(Elaboratable):
    def __init__(self):
        self.count = Signal(64)

    def elaborate(self, platform):
        led = platform.request("led")

        m = Module()

        # Counter
        m.d.sync += self.count.eq(self.count + 1)
        with m.If(self.count >= 50000000):
            m.d.sync += self.count.eq(0)
            m.d.sync += led.eq(~led)

        return m

# To change clock domain of a module:
# new_thing = DomainRenamer("new_clock")(MyElaboratable())

# We sub-class wishbone.Interface here because it needs to be a bus object to be added as a window to Wishbone stuff
class ROM(Elaboratable, Interface):
    def __init__(self, data=None):
        #self.size = len(data)
        self.data = Memory(width=32, depth=(4096 >> 2), init=data)
        self.r = self.data.read_port()

        # Need to init Interface
        Interface.__init__(self, addr_width=10, data_width=32)

        # This is effectively a "window", and it has a certain set of resources
        # 12 = log2(4096)
        memory_map = MemoryMap(addr_width=10, data_width=32)
        # TODO need to unify how I deal with size
        # In this case, one resource, which is out memory
        memory_map.add_resource(self.data, name="rom_data", size=(4096 >> 2))

        self.memory_map = memory_map

    # Connects memory port signals to wishbone interface
    def elaborate(self, platform):
        # Stolen from https://vivonomicon.com/2020/04/14/learning-fpga-design-with-nmigen/
        m = Module()
        # Register the read port submodule.
        m.submodules.r = self.r

        # 'ack' signal should rest at 0.
        m.d.sync += self.ack.eq( 0 )
        # Simulated reads only take one cycle, but only acknowledge
        # them after 'cyc' and 'stb' are asserted.
        with m.If( self.cyc ):
            m.d.sync += self.ack.eq( self.stb )

        # Set 'dat_r' bus signal to the value in the
        # requested 'data' array index.
        m.d.comb += [
            self.r.addr.eq( self.adr ),
            self.dat_r.eq( self.r.data )
        ]

        # End of simulated memory module.
        return m

# TODO support read segmentation or whatever it's called, where you read/write certian bytes from memory
# Otherwise we can't store individual bytes, and this will wreck shit in weird ways.
class RAM(Elaboratable, Interface):
    def __init__(self):
        #self.size = len(data)
        self.data = Memory(width=32, depth=(4096 >> 2))
        self.r = self.data.read_port()
        self.w = self.data.write_port()

        # Need to init Interface
        Interface.__init__(self, addr_width=10, data_width=32)

        # This is effectively a "window", and it has a certain set of resources
        # 12 = log2(4096)
        memory_map = MemoryMap(addr_width=10, data_width=32)
        # TODO need to unify how I deal with size
        # In this case, one resource, which is out memory
        memory_map.add_resource(self.data, name="ram_data", size=(4096 >> 2))

        self.memory_map = memory_map

    # Connects memory port signals to wishbone interface
    def elaborate(self, platform):
        # Stolen from https://vivonomicon.com/2020/04/14/learning-fpga-design-with-nmigen/
        m = Module()
        # Register the read port submodule.
        m.submodules.r = self.r
        m.submodules.w = self.w

        # 'ack' signal should rest at 0.
        m.d.sync += self.ack.eq(0)
        # Simulated reads only take one cycle, but only acknowledge
        # them after 'cyc' and 'stb' are asserted.
        with m.If( self.cyc & self.stb):
            m.d.sync += self.ack.eq(1)

            # Write to address if we are writing
            with m.If(self.we):
                m.d.sync += self.w.en.eq(1)

        # Set 'dat_r' bus signal to the value in the
        # requested 'data' array index.
        m.d.comb += [
            self.r.addr.eq( self.adr ),
            self.dat_r.eq( self.r.data ),
            self.w.addr.eq(self.adr),
            self.w.data.eq(self.dat_w),
        ]

        # End of simulated memory module.
        return m

class LEDPeripheral(Elaboratable, Interface):
    def __init__(self, led_signal):
        Interface.__init__(self, addr_width=1, data_width=32)
        memory_map = MemoryMap(addr_width=1, data_width=32)
        #memory_map.add_resource("my_led", name="led_peripheral", size=1)
        self.memory_map = memory_map

        self.led = led_signal

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

        storage = Signal(1)

        # Always update read values (both wishbone and the LED outpu)
        m.d.comb += [
            self.dat_r[0].eq(storage),
            self.led.eq(storage),
        ]

        m.d.sync += self.ack.eq(0) # default to no ack
        with m.If(self.cyc & self.stb):
            # single cycle ack when CYC and STB are asserted
            m.d.sync += self.ack.eq(1)

            # Write to our storage register if the value has changed
            with m.If(self.we):
                m.d.sync += storage.eq(self.dat_w[0])

        return m


def load_firmware_for_mem() -> List[int]:
    with open('../firmware/fw.bin', 'rb') as f:
        # Stored as little endian, LSB first??
        data = f.read()
        out = []
        assert(len(data) % 4 == 0)
        for i in range(int(len(data) / 4)):
            out.append(int.from_bytes(data[i*4:i*4+4], byteorder='little'))

        return out

class Core(Elaboratable):
    def __init__(self, led_signal):
        self.count = Signal(64)
        self.cpu = Minerva(reset_address=0x01000000)
        self.arbiter = Arbiter(addr_width=32, data_width=32)
        self.decoder = Decoder(addr_width=32, data_width=32)
        self.led_signal = led_signal

    def elaborate(self, platform):
        m = Module()
        m.submodules += self.cpu
        m.submodules += self.arbiter
        m.submodules += self.decoder

        # Connect ibus and dbus together for simplicity for now
        minerva_wb_features = ["cti", "bte", "err"]
        self.ibus = Interface(addr_width=32, data_width=32, features=minerva_wb_features)
        # Note this is a set of statements! without assigning to the comb domain, this will do nothing
        m.d.comb += self.cpu.ibus.connect(self.ibus)
        self.arbiter.add(self.ibus)

        self.dbus = Interface(addr_width=32, data_width=32, features=minerva_wb_features)
        m.d.comb += self.cpu.dbus.connect(self.dbus) # Don't use .eq, use .connect, which will appropriately assign signals
        # using .eq() gave me Multiple Driven errors
        self.arbiter.add(self.dbus)

        # TODO do something with interrupts
        # These are interrupts headed into the CPU
        self.interrupts = Signal(32)
        m.d.comb += [
            # Used for mtime registers, which are memory-mapped (not CSR), so I would have to implement.
            # If I cared.
            self.cpu.timer_interrupt.eq(0),
            # Ostensibly exposed to us so we can interrupt one hart (CPU in this context) from another, we don't need this.
            self.cpu.software_interrupt.eq(0),
            # External interrupt lines, would be for any interrupts I implemented w/ a custom interrupt controller.
            # Most likely I'll map a few lines to some peripherals, won't do anything fancy
            self.cpu.external_interrupt.eq(self.interrupts),
        ]


        fw = load_firmware_for_mem()
        print(len(fw))
        print(fw)

        # Hook up memory space
        self.rom = ROM(fw)
        m.submodules += self.rom
        # Problem: not sure to handle how we do byte vs word addressing properly
        # So doing this shift is a bit of a hacky way to impl anything
        start, _stop, _step = self.decoder.add(self.rom, addr=(0x01000000 >> 2))
        print(f"ROM added at 0x{start:08x}")

        self.ram = RAM()
        m.submodules += self.ram
        start, _stop, _step = self.decoder.add(self.ram)
        print(f"RAM added at 0x{start:08x}")

        self.led = LEDPeripheral(self.led_signal)
        m.submodules += self.led
        start, _stop, _step = self.decoder.add(self.led)
        print(f"LED added at 0x{start:08x}")

        # Connect arbiter to decoder
        m.d.comb += self.arbiter.bus.connect(self.decoder.bus)

        # Counter
        #m.d.sync += self.count.eq(self.count + 1)
        #with m.If(self.count >= 50000000):
        #    m.d.sync += self.count.eq(0)
        #    m.d.sync += led.eq(~led)

        return m

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

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

        led_signal = platform.request("led")
        core = Core(led_signal)
        m.submodules += core

        return m


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


    def elaborate(self, platform):
        m = Module()
        led_signal = Signal()
        core = Core(led_signal)
        m.submodules += core

        return m

def run_sim():
    dut = TestDevice()
    sim = Simulator(dut)

    def proc():
        for i in range(10000):
            yield Tick()

    sim.add_clock(1e-6)
    sim.add_sync_process(proc)
    with sim.write_vcd('test.vcd', gtkw_file='test.gtkw'):
        sim.reset()
        sim.run()


if __name__ == "__main__":
    colorlight_i9.Colorlight_i9_Platform().build(SoC(), debug_verilog=True)

    #run_sim()