#! /usr/bin/python3

import iris.config.resmgr            as resmgr
import iris.config.objects.ring      as ring

import infra.common.objects as objects
from infra.common.logging   import logger
from infra.factory.store    import FactoryStore
import iris.config.objects.rdma.doorbell as doorbell

import model_sim.src.model_wrap as model_wrap


class RdmaRingObject(ring.RingObject):
    def __init__(self):
        super().__init__()
        self.host = None
        self.address = None
        self.size = None
        self.desc_size = None
        self.initialized = False
        self.nic_resident = False
        self.hw_ring_id = 0
        self.qp_spec_en = False

    def Init(self, queue, spec):
        super().Init(queue, spec)
        self.size = spec.size
        self.desc_size = self.descriptor_template.meta.size
        self.doorbell = doorbell.Doorbell()
        self.doorbell.Init(self, spec)

    def SetRingParams(self, ring_id, host, nic_resident, mem_handle, address, size, desc_size):
        self.host = host
        self.nic_resident = nic_resident
        self.address = address
        self.mem_handle = mem_handle
        self.size = size
        self.desc_size = desc_size
        self.hw_ring_id = ring_id
        logger.info("SetRingParams: host %d nic_resident %d base_addr: 0x%x" \
                       " size %d desc_size %d hw_ring_id: %d" %\
                       (self.host, self.nic_resident, self.address, self.size, self.desc_size, self.hw_ring_id))

    def SetRingQpSpecEn(self, spec_en):
        self.qp_spec_en = spec_en

    def Configure(self):
        pass

    def Post(self, descriptor, debug = True):
        #if not self.initialized:
        #    self.queue.qstate.set_ring_base(self.address)
        #    self.queue.qstate.set_ring_size(self.size)
        # Bind the descriptor to the ring
        if debug is True:
            logger.info('posting descriptor at pindex: %d..' %(self.queue.qstate.get_pindex(0)))
        descriptor.address = (self.address + (self.desc_size * self.queue.qstate.get_pindex(0)))
        # pass q_max_desc_size to the descriptor object so that it can check and assert to see 
        # if generated descritptor size is within the limits of q_max_desc_size to avoid
        # memory overwrite
        descriptor.q_max_desc_size = self.desc_size
        if self.nic_resident:
            descriptor.mem_handle = None
        else:
            descriptor.mem_handle = objects.MemHandle(descriptor.address, resmgr.HostMemoryAllocator.v2p(descriptor.address))

        descriptor.Write(self.qp_spec_en)
        logger.info('incrementing pindex..')
        self.queue.qstate.incr_pindex(0, self.size)

        # Increment posted index
        #if self.queue.queue_type.purpose.upper() == "LIF_QUEUE_PURPOSE_RDMA_RECV":
        #    logger.info('incrementing pindex..')
        #    self.queue.qstate.incr_pindex(0)
        # for now, ring doorbell only for SQ.
        # Doorbell ring for RQ will be needed for Prefetch/Cache, will be done later
        #if self.queue.queue_type.purpose == "rdma_sq":
        #    self.doorbell.Ring({})  # HACK

    def Consume(self, descriptor, debug = True):
        if debug is True:
            logger.info("Consuming descriptor on Queue(%s) Ring at cindex: %d" % (self.queue.queue_type.purpose.upper(), self.queue.qstate.get_cindex(0)))
        if self.queue.queue_type.purpose.upper() == "LIF_QUEUE_PURPOSE_RDMA_RECV":
            descriptor.address = (self.address + (self.desc_size * self.queue.qstate.get_proxy_cindex()))
        elif self.queue.queue_type.purpose.upper() == "LIF_QUEUE_PURPOSE_RDMA_SEND":
           descriptor.address = (self.address + (self.desc_size * ((self.queue.qstate.get_cindex(0) - 1) & (self.size - 1))))
        else:
            descriptor.address = (self.address + (self.desc_size * self.queue.qstate.get_cindex(0)))
        if (self.nic_resident):
            descriptor.mem_handle = None 
        else:
            descriptor.mem_handle = objects.MemHandle(descriptor.address, resmgr.HostMemoryAllocator.v2p(descriptor.address))
    
        descriptor.Read()
        self.queue.qstate.Read()

        # Increment consumer index for CQs and EQs
        if ((self.queue.queue_type.purpose.upper() == "LIF_QUEUE_PURPOSE_CQ") or 
            (self.queue.queue_type.purpose.upper() == "LIF_QUEUE_PURPOSE_EQ")):
            logger.info('incrementing cindex..')
            self.queue.qstate.incr_cindex(0, self.size)

    def Read(self):
        raise NotImplementedError


class RdmaRingObjectHelper:
    def __init__(self):
        self.rings = []
        return

    def Generate(self, queue, spec):
        for rspec in spec.rings:
            ring = RdmaRingObject()
            ring.Init(queue, rspec.ring)
            self.rings.append(ring)
        return

    def Configure(self):
        for ring in self.rings:
            ring.Configure()
