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


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


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)
    parser.add_argument('--cores', type=int, default=1)
    args = parser.parse_args()
    return RiscvHWConfig(args.l1i, args.l1d, args.l2,
                         args.vlen, args.elen, args.cores)


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


def createCPU(l2bus, config):
    cpu = RiscvO3CPU()
    cpu.isa = RiscvISA(enable_rvv=True, vlen=config.vlen, elen=config.elen)

    cpu.icache = L1ICache(config)
    cpu.dcache = L1DCache(config)

    cpu.icache.connectCPU(cpu)
    cpu.dcache.connectCPU(cpu)

    cpu.icache.connectBus(l2bus)
    cpu.dcache.connectBus(l2bus)

    cpu.createInterruptController()

    return cpu


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(f"cores: {config.cores}")
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.l2bus = L2XBar()

system.cpu = [createCPU(system.l2bus, config) for _ in range(config.cores)]

system.l2cache = L2Cache(config)
system.l2cache.connectCPUSideBus(system.l2bus)

system.membus = SystemXBar()
system.l2cache.connectMemSideBus(system.membus)

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]
for cpu in system.cpu:
    cpu.workload = process
    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()}")