import m5
from m5.objects import System, SrcClockDomain, VoltageDomain, Root
from m5.objects import RiscvO3CPU, Cache, AddrRange, SEWorkload, Process
from m5.objects import MemCtrl, DDR3_1600_8x8, SystemXBar, L2XBar, RiscvISA


class RiscvHWConfig:
    def __init__(self, l1i, l1d, l2, vlen, elen):
        self.l1i = l1i
        self.l1d = l1d
        self.l2 = l2
        self.vlen = vlen
        self.elen = elen


def get_config():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--l1i', type=str, default='16kB')
    parser.add_argument('--l1d', type=str, default='64kB')
    parser.add_argument('--l2', type=str, default='256kB')
    parser.add_argument('--vlen', type=int, default=256)
    parser.add_argument('--elen', type=int, default=64)
    args = parser.parse_args()
    return RiscvHWConfig(args.l1i, args.l1d, args.l2, args.vlen, args.elen)


class L1Cache(Cache):
    assoc = 2
    tag_latency = 2
    data_latency = 2
    response_latency = 2
    mshrs = 4
    tgts_per_mshr = 20

    def __init__(self, config=None):
        super().__init__()
        pass

    def connectCPU(self, cpu):
        raise NotImplementedError

    def connectBus(self, bus):
        self.mem_side = bus.cpu_side_ports


class L1ICache(L1Cache):
    size = '16kB'

    def __init__(self, config=None):
        super().__init__(config)
        if not config or not config.l1i:
            return
        self.size = config.l1i

    def connectCPU(self, cpu):
        self.cpu_side = cpu.icache_port


class L1DCache(L1Cache):
    size = '64kB'

    def __init__(self, config=None):
        super().__init__(config)
        if not config or not config.l1d:
            return
        self.size = config.l1d

    def connectCPU(self, cpu):
        self.cpu_side = cpu.dcache_port


class L2Cache(Cache):
    size = '256kB'
    assoc = 8
    tag_latency = 20
    data_latency = 20
    response_latency = 20
    mshrs = 20
    tgts_per_mshr = 12

    def __init__(self, config=None):
        super().__init__()
        if not config or not config.l2:
            return
        self.size = config.l2

    def connectCPUSideBus(self, bus):
        self.cpu_side = bus.mem_side_ports

    def connectMemSideBus(self, bus):
        self.mem_side = bus.cpu_side_ports


config = get_config()

print(f"l1i size: {config.l1i}")
print(f"l1d size: {config.l1d}")
print(f"l2 size: {config.l2}")
print(f"vlen size: {config.vlen} bits")
print(f"elen size: {config.elen} bits")
print("\n")

assert config.vlen >= 2 * config.elen, \
    "ERROR: vlen must be at least twice as large as elen"

assert 128 <= config.vlen <= 4096, \
    "ERROR: vlen must be between 128 and 4096 bits"

assert config.vlen & (config.vlen - 1) == 0, \
    "ERROR: vlen must be a power of 2"

assert config.elen in [16, 32, 64], \
    "ERROR: elen must be 16, 32 or 64"

system = System()
system.clk_domain = SrcClockDomain()
system.clk_domain.clock = '1GHz'
system.clk_domain.voltage_domain = VoltageDomain()

system.mem_mode = 'timing'
system.mem_ranges = [AddrRange('512MB')]

system.cpu = RiscvO3CPU()
system.cpu.isa = RiscvISA(enable_rvv=True, vlen=config.vlen, elen=config.elen)

# Create the L1 caches
system.cpu.icache = L1ICache(config)
system.cpu.dcache = L1DCache(config)

# Connect the caches to the CPU
system.cpu.icache.connectCPU(system.cpu)
system.cpu.dcache.connectCPU(system.cpu)

# Connect the CPU to the L2 bus
system.l2bus = L2XBar()

# Connect the L1 caches to the L2 bus
system.cpu.icache.connectBus(system.l2bus)
system.cpu.dcache.connectBus(system.l2bus)

# Connect the L2 cache to the CPU side bus
system.l2cache = L2Cache(config)
system.l2cache.connectCPUSideBus(system.l2bus)

# Connect the L2 cache to the memory bus
system.membus = SystemXBar()
system.l2cache.connectMemSideBus(system.membus)

# Connect the CPU to the memory bus
system.cpu.createInterruptController()

system.mem_ctrl = MemCtrl()
system.mem_ctrl.dram = DDR3_1600_8x8()
system.mem_ctrl.dram.range = system.mem_ranges[0]
system.mem_ctrl.port = system.membus.mem_side_ports

binary = './dct2d_riscv.out'

system.workload = SEWorkload.init_compatible(binary)

process = Process()
process.cmd = [binary]
system.cpu.workload = process
system.cpu.createThreads()

# Run SE mode
root = Root(full_system=False, system=system)
m5.instantiate()

print("Beginning simulation!")
exit_event = m5.simulate()
print(f"Exiting @ tick {m5.curTick()} because {exit_event.getCause()}")