// Copyright (2019-2022) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#include "IBReceiver.h"

#include <sys/mman.h>

#ifdef JFJOCH_USE_NUMA
#include <numaif.h>
#endif

#include "../common/JFJochException.h"
#include "RawJFUDPPacket.h"

#define BUFFER_SIZE 16384
#define BUFFER_COUNT 4096

IBReceiverBuffer::IBReceiverBuffer(uint8_t numa_node) {
    buffer = (uint8_t *) mmap(nullptr, BUFFER_SIZE * BUFFER_COUNT, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (buffer == nullptr)
        throw JFJochException(JFJochExceptionCategory::MemAllocFailed, "frame_buffer");

#ifdef JFJOCH_USE_NUMA
    if (numa_node >= 0) {
        unsigned long nodemask = 1L << numa_node;;
        if (numa_node > sizeof(nodemask)*8) {
            Unmap();
            throw JFJochException(JFJochExceptionCategory::MemAllocFailed, "Mask too small for NUMA node");
        }
        if (mbind(buffer, BUFFER_SIZE * BUFFER_COUNT, MPOL_BIND, &nodemask, sizeof(nodemask)*8, MPOL_MF_STRICT) == -1) {
            Unmap();
            throw JFJochException(JFJochExceptionCategory::MemAllocFailed, "Cannot apply NUMA policy");
        }
    }
#endif
    memset(buffer, 0, BUFFER_SIZE * BUFFER_COUNT);
}

void IBReceiverBuffer::Unmap() {
    munmap(buffer, BUFFER_SIZE * BUFFER_COUNT);
}

IBReceiverBuffer::~IBReceiverBuffer() {
    mr.reset();
    Unmap();
}

void IBReceiverBuffer::Register(IBProtectionDomain &pd) {
    mr = std::make_unique<IBMemoryRegion>(pd, buffer, BUFFER_SIZE * BUFFER_COUNT);
}

IBMemoryRegion *IBReceiverBuffer::GetMemoryRegion() {
    return mr.get();
}

uint8_t *IBReceiverBuffer::GetLocation(uint64_t location) {
    if (location < BUFFER_COUNT)
        return buffer + BUFFER_SIZE * location;
    else
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "Location out of bounds");
}


IBReceiver::IBReceiver(IBContext &context, ProcessJFPacket &process, uint64_t mac_addr, uint32_t ipv4, uint8_t numa_node)
        : pd(context),
          cq(context, BUFFER_COUNT),
          qp(pd, cq, 16, BUFFER_COUNT),
          buffer(numa_node) {
    buffer.Register(pd);

    qp.FlowSteering(mac_addr, ipv4);

    for (int i = 0; i < BUFFER_COUNT - 1; i++)
        qp.PostReceiveWR(*buffer.GetMemoryRegion(), i, buffer.GetLocation(i), BUFFER_SIZE);
    receiver = std::async(std::launch::async, &IBReceiver::Run, this, &process);
}

uint64_t IBReceiver::Run(ProcessJFPacket *process) {
    while (!cancel) {
        int64_t i;
        size_t size;
        if (cq.Poll(i, size) > 0) {
            if (size == sizeof(RawJFUDPacket)) {
                auto ptr = (RawJFUDPacket *) buffer.GetLocation(i);
                process->ProcessPacket(&ptr->jf, ptr->ipv4_header_sour_ip);
                qp.PostReceiveWR(*buffer.GetMemoryRegion(), i, buffer.GetLocation(i), BUFFER_SIZE);
            }
        } else
            std::this_thread::sleep_for(std::chrono::microseconds(10));
    }
    return 0;
}

IBReceiver::~IBReceiver() {
    if (receiver.valid())
        receiver.get();
}

uint64_t IBReceiver::Finalize() {
    cancel = true;
    if (receiver.valid())
        return receiver.get();
    else
        return 0;
}