// Copyright (2019-2023) Paul Scherrer Institute

#ifdef JFJOCH_USE_IBVERBS
#ifndef JUNGFRAUJOCH_IBWRAPPERS_H
#define JUNGFRAUJOCH_IBWRAPPERS_H

#include <infiniband/verbs.h>
#include <string>

class IBContext {
    ibv_context *context = nullptr;
public:
    explicit IBContext(const std::string &dev_name);
    IBContext(const IBContext &context) = delete;
    ~IBContext();
    ibv_port_state GetState();
    uint16_t GetLid();
    ibv_context *GetPointer();
};

class IBProtectionDomain {
    ibv_pd *pd = nullptr;
public:
    explicit IBProtectionDomain(IBContext &context);
    IBProtectionDomain(const IBProtectionDomain &context) = delete;
    ~IBProtectionDomain();
    ibv_pd *GetPointer();
};

class IBCompletionQueue {
    ibv_cq *cq = nullptr;
public:
    IBCompletionQueue(IBContext &context, int minimum_size);
    int Poll(int64_t &id, size_t &compl_size);
    ~IBCompletionQueue();
    ibv_cq *GetPointer();
};


class IBMemoryRegion {
    ibv_mr *buffer_mr = nullptr;
public:
    IBMemoryRegion(IBProtectionDomain &pd, void *pointer, size_t size);
    IBMemoryRegion(const IBMemoryRegion &mr) = delete;
    uint32_t GetLocalKey();
    uint32_t GetRemoteKey();
    ~IBMemoryRegion();
};

class IBQueuePair {
    ibv_qp *qp = nullptr;
    ibv_flow *flow = nullptr;
public:
    IBQueuePair(IBProtectionDomain &pd, IBCompletionQueue &cq, size_t send_queue_size, size_t receive_queue_size);
    uint32_t GetNumber();
    void Init();
    void Reset();
    void ReadyToReceive();
    void ReadyToSend();

    void PostReceiveWR(IBMemoryRegion &mr, uint32_t location, void *pointer, size_t size);
    void PostSendWR(IBMemoryRegion &mr, void *pointer, size_t size);
    void FlowSteeringIPv4(uint32_t ipv4);

    ~IBQueuePair();
};

class IBRegBuffer {
    std::unique_ptr<IBMemoryRegion> mr;
    uint8_t *buffer;
    const size_t loc_count;
    const size_t loc_size;
    void Unmap();
public:
    explicit IBRegBuffer(IBProtectionDomain &pd, size_t loc_count, size_t loc_size, int16_t numa_node = -1);
    ~IBRegBuffer();
    uint8_t *GetLocation(uint64_t location);
    IBMemoryRegion *GetMemoryRegion();
};

#endif //JUNGFRAUJOCH_IBWRAPPERS_H
#endif // JFJOCH_USE_IBVERBS