#!/usr/bin/env python3

from amaranth import *
from amaranth.sim import *
from platforms import *
from amaranth_soc.wishbone import *
from amaranth_soc import csr
from amaranth_soc.csr.wishbone import WishboneCSRBridge
from amaranth.lib.io import pin_layout

from minerva.core import Minerva

from typing import List
from argparse import ArgumentParser
from pathlib import Path
import unittest
import sys
import os

from memory import *
from led import *
from timer import *
from eth import *

import i2c
import test_i2c
import uart

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


# TODO clean this up, generate binary here, maybe even run cargo build
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 I2CPads:
    """Small structure to avoid using record. Unsure what the proper solution is"""
    class Pad:
        def __init__(self):
            self.o = Signal()
            self.i = Signal()
            self.oe = Signal()

    def __init__(self):
        self.scl = self.Pad()
        self.sda = self.Pad()

class Core(Elaboratable):
    def __init__(self, clk25, led_signal, eth_interface):
        self.count = Signal(64)
        self.cpu = Minerva(reset_address=0x01000000, with_debug=False)
        self.arbiter = Arbiter(addr_width=30, data_width=32, granularity=8)
        self.decoder = Decoder(addr_width=30, data_width=32, granularity=8, features=["err"])
        self.clk25 = clk25
        self.led_signal = led_signal
        self.eth_interface = eth_interface

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

        # Create main and sampling clock, using PLL and 25MHz input clock
        if platform is not None:
            platform.add_file("pll.v", open("pll.v", "r").read())
        sys_clk = Signal()
        sample_clk = Signal()
        pll = Instance(
            "pll",
            i_clkin=self.clk25,
            o_clkout0=sys_clk,
            o_clkout1=sample_clk,
        )
        m.submodules.pll = pll

        # Create new clock domains
        m.domains += ClockDomain("sync")
        m.domains += ClockDomain("sample")
        m.d.comb += ClockSignal("sync").eq(sys_clk)
        m.d.comb += ClockSignal("sample").eq(sample_clk)

        # Add clock constraints
        if platform is not None:
            platform.add_clock_constraint(sys_clk, 50e6)
            platform.add_clock_constraint(sample_clk, 10e6)

        # Connect ibus and dbus together for simplicity for now
        minerva_wb_features = ["cti", "bte", "err"]
        self.ibus = Interface(addr_width=30, data_width=32, granularity=8, 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=30, data_width=32, granularity=8, 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()

        # Hook up memory space
        self.rom = ROM(size_bytes=1024 * 64, data=fw)
        m.submodules.rom = 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)
        print(f"ROM added at 0x{start:08x}")

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

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

        self.timer = TimerPeripheral(50e6, 1e3)
        m.submodules.timer = self.timer
        start, _stop, _step = self.decoder.add(self.timer, addr=0x01300000)
        print(f"Timer added at 0x{start:08x}")

        # TODO how to set addr_width?
        # Create CSR bus and connect it to Wishbone
        self.csr = csr.Decoder(addr_width=10, data_width=8)
        m.submodules.csr = self.csr

        # I2C (connected to DAC for VCO and ADC?)
        if platform is not None:
            i2c_pads = platform.request("i2c")
        else:
            # TODO this is hacky crap
            i2c_pads = I2CPads()
            with m.If(i2c_pads.scl.oe):
                m.d.comb += i2c_pads.scl.i.eq(0)
            with m.Else():
                m.d.comb += i2c_pads.scl.i.eq(1)
        self.i2c = i2c.I2C(50e6, 100e3, i2c_pads)
        m.submodules.i2c = self.i2c
        i2c_start, _stop, _step = self.csr.add(self.i2c.bus)
        print(f"LED added to CSR at 0x{i2c_start}")

        if platform is not None:
            uart_pads = platform.request("uart", 1)
        else:
            uart_pads = None
        # TODO spread sysclk freq through design
        self.uart = uart.UART(50e6, 115_200, pins=uart_pads)
        m.submodules.uart = self.uart
        uart_start, _stop, _step = self.csr.add(self.uart.bus)
        print(f"UART added to CSR at 0x{uart_start:x}")

        self.csr_bridge = WishboneCSRBridge(self.csr.bus, data_width=32, name="CSR")
        m.submodules.csr_bridge = self.csr_bridge
        # TODO shouldn't have to hard-specify this address
        start, _stop, _step = self.decoder.add(self.csr_bridge.wb_bus, addr=0x02000000)
        print(f"CSR bus added at 0x{start:08x}")

        # Ethernet
        self.eth = LiteEth(self.eth_interface)
        m.submodules.eth = self.eth
        start, _stop, _step = self.decoder.add(self.eth, addr=0x03000000)
        print(f"LiteETH 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):
        # TODO pull I2C and UART into here instead of the "core"
        if platform is not None:
            clk25 = platform.request("clk25")
            led_signal = platform.request("led")
            ethernet_interface = platform.request("eth_rgmii", 1)
        else:
            # platform is None in simulation, so provide harnesses for required signals
            clk25 = Signal() # TODO unsure if this will work in sim
            led_signal = Signal()
            ethernet_interface = Record(rgmii_layout)

        return Core(clk25, led_signal, ethernet_interface)


# TODO add structure to add regression tests
def run_sim():
    dut = SoC()
    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('sim.vcd', gtkw_file='sim.gtkw'):
        sim.reset()
        sim.run()


if __name__ == "__main__":
    args = ArgumentParser(description="ARVP Sonar Acquisition FPGA gateware.")
    args.add_argument("--build", action="store_true", help="Build bitstream.")
    args.add_argument("--gen-debug-verilog", action="store_true", help="Save debug verilog.")
    # TODO maybe allow an optional arg to specify an individual test to run?
    args.add_argument("--test", action="store_true", help="Run RTL test suite and report results.")
    args.add_argument("--save-vcd", action="store_true", help="Save VCD waveforms from test(s).")
    args.add_argument("--sim", action="store_true", help="Run overall simulation")
    args = args.parse_args()

    if args.build:
        # Overrides are available via AMARANTH_<override_variable_name> env variable, or kwarg
        # TODO fix platform so I don't have to manually specify MDIO signal
        Colorlight_i9_Platform().build(SoC(), debug_verilog=args.gen_debug_verilog, nextpnr_opts="--router router2")

    if args.test:
        if args.save_vcd:
            if not Path("vcd_out").exists():
                os.mkdir("vcd_out")

            os.environ["TEST_SAVE_VCD"] = "YES"

        # Super hacky... why am I doing this
        test_modules = [mod for mod in sys.modules if mod.startswith("test_")]
        for mod in test_modules:
            unittest.main(module=mod, argv=[sys.argv[0]])

    if args.sim:
        print("Running simulation...")
        run_sim()