import argparse
import struct

from nmigen import *
from nmigen.back import rtlil, pysim

from heavycomps import uart, wishbone
from minerva.core import Minerva


class SimpleWishboneSerial(Elaboratable):
    def __init__(self, tx, sys_clk_freq, baudrate=115200):
        self.tx = tx
        self.bus = wishbone.Interface()
        self.ftw = round(2**32*baudrate/sys_clk_freq)

    def elaborate(self, platform):
        m = Module()
        m.submodules.tx = tx = uart.RS232TX(self.ftw)
        m.d.comb += [
            tx.stb.eq(self.bus.cyc & self.bus.stb & self.bus.we),
            tx.data.eq(self.bus.dat_w),
            self.bus.ack.eq(tx.ack),
            self.tx.eq(tx.tx)
        ]
        return m
        

class Top(Elaboratable):
    def __init__(self, firmware, create_clock):
        if create_clock:
            self.clk100 = Signal()
        self.led = Signal()
        self.serial_tx = Signal()
        self.firmware = firmware

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

        if hasattr(self, "clk100"):
            cd_sync = ClockDomain(reset_less=True)
            m.domains += cd_sync
            m.d.comb += cd_sync.clk.eq(self.clk100)

        counter = Signal(27)
        m.d.sync += counter.eq(counter + 1)
        m.d.comb += self.led.eq(counter[-1])

        m.submodules.cpu = cpu = Minerva(with_icache=False, with_dcache=False, with_muldiv=False)
        m.submodules.ram = ram = wishbone.SRAM(Memory(32, 1024, init=self.firmware))
        m.submodules.uart = uart = SimpleWishboneSerial(self.serial_tx, 100e6)
        m.submodules.con = con = wishbone.InterconnectShared(
            [cpu.ibus, cpu.dbus],
            [
                (lambda a: ~a[20], ram.bus),
                (lambda a:  a[20], uart.bus)
            ], register=True)

        return m


def read_firmware(file):
    firmware = []
    with open(file, "rb") as f:
        while True:
            word = f.read(4)
            if len(word) < 4:
                break
            firmware.append(struct.unpack("<I", word)[0])
    return firmware


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--simulate", action="store_true")
    parser.add_argument("firmware_bin")
    parser.add_argument("output_file")
    args = parser.parse_args()

    firmware = read_firmware(args.firmware_bin)
    top = Top(firmware, create_clock=not args.simulate)

    if args.simulate:
        with pysim.Simulator(top,
                vcd_file=open(args.output_file + ".vcd", "w"),
                gtkw_file=open(args.output_file + ".gtkw", "w")) as sim:
            sim.add_clock(1e-6)
            sim.run_until(1000e-6, run_passive=True)
    else:
        output = rtlil.convert(Fragment.get(top, None),
            ports=(top.clk100, top.led, top.serial_tx))
        with open(args.output_file, "w") as f:
            f.write(output)

if __name__ == "__main__":
    main()