/*
 *  $ gcc -g -Wall -libverbs -o ibverbs-sample1 ibverbs-sample1.c
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <assert.h>
#include <inttypes.h>
#include <infiniband/verbs.h>

enum {
    PORT_NUM = 1,
};

char buffer[1024 * 1024];

static void test_run(struct ibv_context *context, uint16_t lid);
static void modify_qp(struct ibv_qp *qp, uint32_t src_psn, uint16_t dest_lid, uint32_t dest_pqn, uint32_t dest_psn);
static void post_recv(struct ibv_qp *qp, struct ibv_mr *mr, void *addr, uint32_t length);
static void post_send(struct ibv_qp *qp, struct ibv_mr *mr, void *addr, uint32_t length);

int main(int argc, char **argv)
{
    int ret;

    ret = ibv_fork_init();
    if (ret) {
        fprintf(stderr, "Failure: ibv_fork_init (errno=%d)\n", ret);
        exit(EXIT_FAILURE);
    }

    struct ibv_device **dev_list;
    dev_list = ibv_get_device_list(NULL);

    if (!dev_list) {
        int errsave = errno;
        fprintf(stderr, "Failure: ibv_get_device_list (errno=%d)\n", errsave);
        exit(EXIT_FAILURE);        
    }

    if (dev_list[0]) {
        struct ibv_device *device = dev_list[0];

        printf("IB device: %s\n", ibv_get_device_name(device));

        struct ibv_context *context;
        context = ibv_open_device(device);
        assert(context);

        struct ibv_port_attr port_attr;
        ret = ibv_query_port(context, PORT_NUM, &port_attr);
        assert(ret == 0);
        assert(port_attr.lid != 0);

        test_run(context, port_attr.lid);

        ibv_close_device(context);
    }

    ibv_free_device_list (dev_list);

    return 0;
}

static void
test_run(struct ibv_context *context, uint16_t lid)
{
    int i, ret;
    struct ibv_pd *pd;
    struct ibv_mr *mr;
    struct ibv_cq *cq;
    struct ibv_qp *qp1, *qp2;

    pd = ibv_alloc_pd(context);
    assert(pd);

    mr = ibv_reg_mr(pd, buffer, sizeof(buffer), IBV_ACCESS_LOCAL_WRITE);
    assert(mr);

    cq = ibv_create_cq(context, 64, NULL, NULL, 0);
    assert(cq);

    struct ibv_qp_init_attr qp_init_attr = {
        .qp_type    = IBV_QPT_RC,
        .send_cq    = cq,
        .recv_cq    = cq,
        .cap        = {
            .max_send_wr  = 32,
            .max_recv_wr  = 32,
            .max_send_sge =  1,
            .max_recv_sge =  1,
        },
        .sq_sig_all = 1, 
    };

    qp1 = ibv_create_qp(pd, &qp_init_attr);
    assert(qp1);
    qp2 = ibv_create_qp(pd, &qp_init_attr);
    assert(qp2);

    uint32_t psn1 = rand() & 0xFFFFFF;
    uint32_t psn2 = rand() & 0xFFFFFF;

    modify_qp(qp1, psn1, lid, qp2->qp_num, psn2);
    modify_qp(qp2, psn2, lid, qp1->qp_num, psn1);
    
    // post receive work request
    post_recv(qp1, mr, buffer,        4096);

    // post send work request
    for (i=0 ; i<1000 ; i++)
        buffer[4096 + i] = rand();
    post_send(qp2, mr, buffer + 4096, 1000);

    // poll cq
    int num_wr = 2;
    struct ibv_wc wc;

    while (num_wr > 0) {
        ret = ibv_poll_cq(cq, 1, &wc);

        if (ret == 0)
            continue; /* polling */

        if (ret < 0) {
            fprintf(stderr, "Failure: ibv_poll_cq\n");
            exit(EXIT_FAILURE);
        }
        
        if (wc.status != IBV_WC_SUCCESS) {
            fprintf(stderr, "Completion errror\n");
            exit(EXIT_FAILURE);
        }
    
        switch (wc.opcode) {
        case IBV_WC_SEND:
            printf("poll send wc: wr_id=0x%016" PRIx64 "\n", wc.wr_id);
            break;

        case IBV_WC_RECV:
            printf("poll recv wc: wr_id=0x%016" PRIx64 " byte_len=%u, imm_data=0x%x\n", wc.wr_id, wc.byte_len, wc.imm_data);
            // 送信バッファと受信バッファの内容を比較
            ret = memcmp(buffer + 4096, (char*)(uintptr_t)wc.wr_id, wc.byte_len);
            assert(ret == 0);
            break;

        default:
            exit(EXIT_FAILURE);
        }

        num_wr--;        
    }

    ibv_destroy_qp(qp2);
    ibv_destroy_qp(qp1);
    ibv_destroy_cq(cq);
    ibv_dereg_mr(mr);
    ibv_dealloc_pd(pd);

    printf("OK\n");
}

static void
modify_qp(struct ibv_qp *qp, uint32_t src_psn, uint16_t dest_lid, uint32_t dest_pqn, uint32_t dest_psn)
{
    int ret;

    struct ibv_qp_attr init_attr = {
        .qp_state        = IBV_QPS_INIT,
        .port_num        = PORT_NUM,
        .qp_access_flags = IBV_ACCESS_LOCAL_WRITE,
    };

    ret = ibv_modify_qp(qp, &init_attr,
                        IBV_QP_STATE|IBV_QP_PKEY_INDEX|IBV_QP_PORT|IBV_QP_ACCESS_FLAGS);
    assert(ret == 0);

    struct ibv_qp_attr rtr_attr = {
        .qp_state               = IBV_QPS_RTR,
        .path_mtu               = IBV_MTU_4096,
        .dest_qp_num            = dest_pqn,
        .rq_psn                 = dest_psn,
        .max_dest_rd_atomic     = 0,
        .min_rnr_timer          = 0,
        .ah_attr                = {
            .is_global          = 0,
            .dlid               = dest_lid,
            .sl                 = 0,
            .src_path_bits      = 0,
            .port_num           = PORT_NUM,
        },
    };

    ret = ibv_modify_qp(qp, &rtr_attr,
                        IBV_QP_STATE|IBV_QP_AV|IBV_QP_PATH_MTU|IBV_QP_DEST_QPN|
                        IBV_QP_RQ_PSN|IBV_QP_MAX_DEST_RD_ATOMIC|IBV_QP_MIN_RNR_TIMER);
    assert(ret == 0);

    struct ibv_qp_attr rts_attr = {
        .qp_state           = IBV_QPS_RTS,
        .timeout            = 0,
        .retry_cnt          = 7,
        .rnr_retry          = 7,
        .sq_psn             = src_psn,
        .max_rd_atomic      = 0,
    };

    ret = ibv_modify_qp(qp, &rts_attr,
                        IBV_QP_STATE|IBV_QP_TIMEOUT|IBV_QP_RETRY_CNT|IBV_QP_RNR_RETRY|IBV_QP_SQ_PSN|IBV_QP_MAX_QP_RD_ATOMIC);
    assert(ret == 0);
}

static void
post_recv(struct ibv_qp *qp, struct ibv_mr *mr, void *addr, uint32_t length)
{
    int ret;

    struct ibv_sge sge = {
        .addr   = (uint64_t)(uintptr_t)addr,
        .length = length,
        .lkey   = mr->lkey,
    };

    struct ibv_recv_wr recv_wr = {
        .wr_id   = (uint64_t)(uintptr_t)addr,
        .sg_list = &sge,
        .num_sge = 1,
    };

    struct ibv_recv_wr *bad_wr;
    ret = ibv_post_recv(qp, &recv_wr, &bad_wr);
    assert(ret == 0);

    printf("post recv wr\n");
}

static void
post_send(struct ibv_qp *qp, struct ibv_mr *mr, void *addr, uint32_t length)
{
    int ret;
    struct ibv_sge sge = {
        .addr   = (uint64_t)(uintptr_t)addr,
        .length = length,
        .lkey   = mr->lkey,
    };

    struct ibv_send_wr send_wr = {
        .wr_id      = (uint64_t)(uintptr_t)addr,
        .sg_list    = &sge,
        .num_sge    = 1,
        .opcode     = IBV_WR_SEND_WITH_IMM,
        .imm_data   = rand(),
    };

    struct ibv_send_wr *bad_wr;
    ret = ibv_post_send(qp, &send_wr, &bad_wr);
    assert(ret == 0);

    printf("post send wr: imm_data=0x%08x, byte_len=%u\n", send_wr.imm_data, length);
}