#include <stdarg.h>
#include <stddef.h>
#include <stdint.h>
#include <setjmp.h>
#include <string.h>
#include <cmocka.h>
#include <marsio.h>

#include <rte_mbuf.h>
#include <rte_mempool.h>

struct rte_mempool * mbuf_pool = NULL;

int __wrap_vnode_mirror_dequeue_burst(void * cons, unsigned int qid, void * buffs[], unsigned int nr_buffs)
{
    void ** mock_buffs = mock_ptr_type(void **);
    unsigned int mock_nr_buffs = mock_type(unsigned int);
    unsigned int mock_expect_qid = mock_type(unsigned int);

    if (qid == mock_expect_qid)
    {
        for (unsigned int i = 0; i < mock_nr_buffs; i++)
        {
            buffs[i] = mock_buffs[i];
        }
        return (int)mock_nr_buffs;
    }
    else
    {
        return 0;
    }
}

int __wrap_vnode_mirror_enqueue_bulk(void * prod, unsigned int prodq, struct rte_mbuf * buffs[],
                              uint32_t hash[], unsigned int nr_buffs)
{
    /* check input parameters */
    check_expected_ptr(buffs);
    check_expected(nr_buffs);

    /* fault inject parameter */
    int enq_succ_buff_count = mock_type(int);
    assert(nr_buffs >= enq_succ_buff_count);

    /* free the succ enq pkts */
    for (unsigned int i = 0; i < enq_succ_buff_count; i++)
    {
        rte_pktmbuf_free(buffs[i]);
    }

    return 0;
}

int __wrap_vnode_mirror_rt_object_retrieve(void * prod, unsigned int prodq, struct rte_mbuf * rt_objs[],
                                           unsigned int nr_max_rt_objects)
{

    void ** mock_buffs = mock_ptr_type(void **);
    unsigned int mock_nr_buffs = mock_type(unsigned int);

    for (unsigned int i = 0; i < mock_nr_buffs; i++)
    {
        rt_objs[i] = mock_buffs[i];
    }

    return (int)mock_nr_buffs;
}

int __wrap_marsio_buff_free_v2(struct mr_instance * instance, marsio_buff_t * buffs[], unsigned int nr_buffs,
                               unsigned int flags, unsigned int sid)
{
    for (unsigned int i = 0; i < nr_buffs; i++)
    {
        rte_pktmbuf_free(buffs[i]);
    }

    return 0;
}

void __wrap_marsio_buff_do_rehash(struct mr_instance * mr_instance, marsio_buff_t * m)
{
    function_called();
}

int __wrap_marsio_dp_trace_measurements_can_emit(struct mr_instance * instance, const marsio_buff_t * mbuf,
                                                                    uint8_t measurement_type)
{
    return 0;
}

#include <mrapp.h>
struct mr_vdev * __wrap_marsio_open_device(struct mr_instance * instance, const char * devsym,
                                                                unsigned int nr_rxstream,
                                                                unsigned int nr_txstream)
{
    /* create a fake mr_vdev handle */
    struct mr_vdev * mock_mr_vdev_handle = malloc(sizeof(struct mr_vdev));
    memset(mock_mr_vdev_handle, 0, sizeof(struct mr_vdev));

    mock_mr_vdev_handle->nr_rxstream = nr_rxstream;
    mock_mr_vdev_handle->nr_txstream = nr_txstream;
    mock_mr_vdev_handle->instance = instance;

    mock_mr_vdev_handle->vdi = calloc(1, sizeof(struct vdev_instance));
    mock_mr_vdev_handle->tap_representor = NULL;

#define SZ_VDEV_BUFFER 32
    for (unsigned int i = 0; i < mock_mr_vdev_handle->nr_rxstream; i++)
    {
        size_t sz_rx_buffer = sizeof(struct mr_vdev_rx_buffer) + sizeof(struct rte_mbuf *) * SZ_VDEV_BUFFER;
        mock_mr_vdev_handle->rx_buffer[i] = calloc(1, sz_rx_buffer);
        mock_mr_vdev_handle->rx_buffer[i]->size = SZ_VDEV_BUFFER;
    }

    for (unsigned int i = 0; i < mock_mr_vdev_handle->nr_txstream; i++)
    {
        size_t sz_tx_buffer = sizeof(struct mr_vdev_tx_buffer) + sizeof(struct rte_mbuf *) * SZ_VDEV_BUFFER;
        mock_mr_vdev_handle->tx_buffer[i] = calloc(1, sz_tx_buffer);
        mock_mr_vdev_handle->tx_buffer[i]->size = SZ_VDEV_BUFFER;
    }

    assert_non_null(mock_mr_vdev_handle);
    return mock_mr_vdev_handle;
}

void __wrap_marsio_close_device(struct mr_vdev * vdev)
{
    for (unsigned int i = 0; i < vdev->nr_rxstream; i++)
    {
        free(vdev->rx_buffer[i]);
    }

    for (unsigned int i = 0; i < vdev->nr_txstream; i++)
    {
        free(vdev->tx_buffer[i]);
    }

    free(vdev->vdi);
    free(vdev);
}

static void test_recv_burst_api_one_pkt(void ** state)
{
    struct mr_vdev * mock_mr_vdev = marsio_open_device(NULL,"mock_mr_vdev", 2, 1);
    assert_non_null(mock_mr_vdev);

    struct rte_mbuf * mbufs_to_recv[1];
    mbufs_to_recv[0] = rte_pktmbuf_alloc(mbuf_pool);
    assert_non_null(mbufs_to_recv[0]);

    will_return(__wrap_vnode_mirror_dequeue_burst, mbufs_to_recv);
    will_return(__wrap_vnode_mirror_dequeue_burst, 1);
    will_return(__wrap_vnode_mirror_dequeue_burst, 0);

    struct rte_mbuf * mbufs_to_recv_ret[1];
    int ret = marsio_recv_burst(mock_mr_vdev, 0, (marsio_buff_t **)mbufs_to_recv_ret, 1);
    assert_int_equal(ret, 1);

    /* free the mbuf */
    assert_true(mbufs_to_recv_ret[0] == mbufs_to_recv[0]);
    rte_pktmbuf_free(mbufs_to_recv[0]);

    /* make sure no mbuf leak */
    assert_int_equal(rte_mempool_in_use_count(mbuf_pool), 0);

    will_return(__wrap_vnode_mirror_dequeue_burst, NULL);
    will_return(__wrap_vnode_mirror_dequeue_burst, 1);
    will_return(__wrap_vnode_mirror_dequeue_burst, 0);

    /* try to recv pkt from another queue */
    ret = marsio_recv_burst(mock_mr_vdev, 1, (marsio_buff_t **)mbufs_to_recv_ret, 1);
    assert_int_equal(ret, 0);

    /* try to recv pkt from out-of-range queue */
    ret = marsio_recv_burst(mock_mr_vdev, 2, (marsio_buff_t **)mbufs_to_recv_ret, 1);
    assert_int_equal(ret, -EINVAL);
}

static void test_recv_burst_api_many_pkts(void ** state)
{
    struct rte_mbuf * mbufs_to_recv[32] = {};
    for (unsigned int i = 0; i < 32; i++)
    {
        mbufs_to_recv[i] = rte_pktmbuf_alloc(mbuf_pool);
        assert_non_null(mbufs_to_recv[i]);
    }

    /* dequeue will get 32 mbufs */
    will_return(__wrap_vnode_mirror_dequeue_burst, mbufs_to_recv);
    will_return(__wrap_vnode_mirror_dequeue_burst, 32);
    will_return(__wrap_vnode_mirror_dequeue_burst, 0);

    /* call marsio_recv_burst one by one, should get 32 mbufs */
    struct mr_vdev * mock_mr_vdev = marsio_open_device(NULL, "mock_mr_vdev", 1, 1);
    for (unsigned int i = 0; i < 32; i++)
    {
        struct rte_mbuf * mbufs_to_recv_ret[1];
        int ret = marsio_recv_burst(mock_mr_vdev, 0, (marsio_buff_t **)mbufs_to_recv_ret, 1);
        assert_int_equal(ret, 1);

        /* free the mbuf */
        assert_true(mbufs_to_recv_ret[0] == mbufs_to_recv[i]);
        rte_pktmbuf_free(mbufs_to_recv[i]);
    }

    /* make sure no mbuf leak */
    assert_int_equal(rte_mempool_in_use_count(mbuf_pool), 0);
}

static void test_send_burst_api_one_pkt_succ(void ** state)
{
    struct mr_vdev * mock_mr_vdev = marsio_open_device(NULL, "mock_mr_vdev", 1, 1);
    assert_non_null(mock_mr_vdev);

    struct mr_sendpath * mock_sendpath = marsio_sendpath_create_by_vdev(mock_mr_vdev);
    assert_non_null(mock_sendpath);

    /* create a mbuf to send */
    struct rte_mbuf * mbufs_to_send[1];
    mbufs_to_send[0] = rte_pktmbuf_alloc(mbuf_pool);
    assert_non_null(mbufs_to_send[0]);

    /* no packet drop because of backpressure, all the packets should be sent */
    will_return(__wrap_vnode_mirror_rt_object_retrieve, NULL);
    will_return(__wrap_vnode_mirror_rt_object_retrieve, 0);

    /* mbufs_to_send should be call as parameter for vnode_mirror_enqueue_bulk */
    expect_memory(__wrap_vnode_mirror_enqueue_bulk, buffs, mbufs_to_send, 1);
    expect_value(__wrap_vnode_mirror_enqueue_bulk, nr_buffs, 1);
    will_return(__wrap_vnode_mirror_enqueue_bulk, RTE_DIM(mbufs_to_send));

    int ret = marsio_send_burst(mock_sendpath, 0, (marsio_buff_t **)mbufs_to_send, 1);
    assert_int_equal(ret, 0);

    /* flush */
    struct mr_vdev * vdevs_handles[] = {mock_mr_vdev};
    marsio_poll_wait(mock_mr_vdev->instance, vdevs_handles, RTE_DIM(vdevs_handles), 0, 0);

    /* make sure no mbuf leak */
    assert_true(rte_mempool_full(mbuf_pool));
    marsio_close_device(mock_mr_vdev);
}

static void test_send_burst_api_one_pkt_drop_by_backpressure(void ** state)
{
    struct mr_vdev * mock_mr_vdev = marsio_open_device(NULL, "mock_mr_vdev", 1, 1);
    assert_non_null(mock_mr_vdev);

    struct mr_sendpath * mock_sendpath = marsio_sendpath_create_by_vdev(mock_mr_vdev);
    assert_non_null(mock_sendpath);

    /* create a mbuf to send */
    struct rte_mbuf * mbufs_to_send[1];
    mbufs_to_send[0] = rte_pktmbuf_alloc(mbuf_pool);
    assert_non_null(mbufs_to_send[0]);

    /* mbufs_to_send should be call as parameter for vnode_mirror_enqueue_bulk */
    expect_memory(__wrap_vnode_mirror_enqueue_bulk, buffs, mbufs_to_send, 1);
    expect_value(__wrap_vnode_mirror_enqueue_bulk, nr_buffs, 1);

    /* fault inject, no packet send successfully */
    will_return(__wrap_vnode_mirror_enqueue_bulk, 0);

    /* all packet should be retrieved to simulate the backpressure */
    will_return(__wrap_vnode_mirror_rt_object_retrieve, mbufs_to_send);
    will_return(__wrap_vnode_mirror_rt_object_retrieve, 1);

    int ret = marsio_send_burst(mock_sendpath, 0, (marsio_buff_t **)mbufs_to_send, 1);
    assert_int_equal(ret, 0);

    /* flush */
    struct mr_vdev * vdevs_handles[] = {mock_mr_vdev};
    marsio_poll_wait(mock_mr_vdev->instance, vdevs_handles, RTE_DIM(vdevs_handles), 0, 0);

    /* make sure no mbuf leak */
    assert_true(rte_mempool_full(mbuf_pool));
    marsio_close_device(mock_mr_vdev);
}

struct send_with_options_test_state
{
    struct mr_vdev * mock_mr_vdev;
    struct mr_sendpath * mock_sendpath;
};

/* test group 2: send with options */
static int testcase_group_SEND_WITH_OPTIONS_setup(void ** state)
{
    struct send_with_options_test_state * test_state = malloc(sizeof(struct send_with_options_test_state));
    memset(test_state, 0, sizeof(struct send_with_options_test_state));

    test_state->mock_mr_vdev = marsio_open_device(NULL, "mock_mr_vdev", 1, 1);
    assert_non_null(test_state->mock_mr_vdev);

    test_state->mock_sendpath = marsio_sendpath_create_by_vdev(test_state->mock_mr_vdev);
    assert_non_null(test_state->mock_sendpath);

    *state = test_state;
    return 0;
}

static int testcase_group_SEND_WITH_OPTIONS_teardown(void ** state)
{
    struct send_with_options_test_state * test_state = *state;
    if (test_state != NULL)
    {
        marsio_close_device(test_state->mock_mr_vdev);
        marsio_sendpath_destory(test_state->mock_sendpath);
        free(test_state);
    }

    return 0;
}

static void test_send_with_options_do_not_free(void ** state)
{
    struct send_with_options_test_state * test_state = *state;
    assert_non_null(test_state);

    struct rte_mbuf * mbufs_to_send[1];
    mbufs_to_send[0] = rte_pktmbuf_alloc(mbuf_pool);
    assert_non_null(mbufs_to_send[0]);

    /* mbufs_to_send should be call as parameter for vnode_mirror_enqueue_bulk */
    expect_memory(__wrap_vnode_mirror_enqueue_bulk, buffs, mbufs_to_send, 1);
    expect_value(__wrap_vnode_mirror_enqueue_bulk, nr_buffs, 1);

    /* packet should be sent successfully, no packet drop */
    will_return(__wrap_vnode_mirror_enqueue_bulk, 1);
    will_return(__wrap_vnode_mirror_rt_object_retrieve, NULL);
    will_return(__wrap_vnode_mirror_rt_object_retrieve, 0);

    /* call send with option, do not free the packet */
    int ret = marsio_send_burst_with_options(test_state->mock_sendpath, 0, (marsio_buff_t **)mbufs_to_send, 1,
                                             MARSIO_SEND_OPT_NO_FREE);

    /* always success */
    assert_int_equal(ret, 0);

    /* since the pkt has been set to no free */
    assert_int_equal(rte_mempool_in_use_count(mbuf_pool), 1);

    /* free the mbuf */
    rte_pktmbuf_free(mbufs_to_send[0]);
    assert_true(rte_mempool_full(mbuf_pool));
}

static void test_send_with_options_rehash(void ** state)
{
    struct send_with_options_test_state * test_state = *state;
    assert_non_null(test_state);

    struct rte_mbuf * mbufs_to_send[1];
    mbufs_to_send[0] = rte_pktmbuf_alloc(mbuf_pool);

    /* mbufs_to_send should be call as parameter for vnode_mirror_enqueue_bulk */
    expect_memory(__wrap_vnode_mirror_enqueue_bulk, buffs, mbufs_to_send, 1);
    expect_value(__wrap_vnode_mirror_enqueue_bulk, nr_buffs, 1);

    /* packet should be sent successfully, no packet drop */
    will_return(__wrap_vnode_mirror_enqueue_bulk, 1);
    will_return(__wrap_vnode_mirror_rt_object_retrieve, NULL);
    will_return(__wrap_vnode_mirror_rt_object_retrieve, 0);

    /* marsio_do_rehash should be called */
    expect_function_call(__wrap_marsio_buff_do_rehash);

    /* call send with option, do rehash */
    int ret = marsio_send_burst_with_options(test_state->mock_sendpath, 0, (marsio_buff_t **)mbufs_to_send, 1,
                                             MARSIO_SEND_OPT_REHASH);

    /* always success */
    assert_int_equal(ret, 0);

    /* check the mbuf leak */
    assert_true(rte_mempool_full(mbuf_pool));
}

static void test_send_with_options_many_pkts(void ** state)
{
    struct send_with_options_test_state * test_state = *state;
    assert_non_null(test_state);

    struct rte_mbuf * mbufs_to_send[32] = {};
    for (unsigned int i = 0; i < 32; i++)
    {
        mbufs_to_send[i] = rte_pktmbuf_alloc(mbuf_pool);
        assert_non_null(mbufs_to_send[i]);
    }

    /* mbufs_to_send should be call as parameter for vnode_mirror_enqueue_bulk */
    expect_memory(__wrap_vnode_mirror_enqueue_bulk, buffs, mbufs_to_send, 32);
    expect_value(__wrap_vnode_mirror_enqueue_bulk, nr_buffs, 32);

    /* fault inject, all packet send successfully */
    will_return(__wrap_vnode_mirror_enqueue_bulk, 32);
    will_return(__wrap_vnode_mirror_rt_object_retrieve, NULL);
    will_return(__wrap_vnode_mirror_rt_object_retrieve, 0);

    /* call send with option, no options */
    int ret = marsio_send_burst_with_options(test_state->mock_sendpath, 0, (marsio_buff_t **)mbufs_to_send, 32, 0);
    assert_int_equal(ret, 0);

    /* make sure no mbuf leak */
    assert_int_equal(rte_mempool_in_use_count(mbuf_pool), 0);
}

#if 0
static void test_mix_send_with_option_and_send_burst(void ** state)
{
    struct send_with_options_test_state * test_state = *state;
    assert_non_null(test_state);

    struct rte_mbuf * mbufs_to_send[32] = {};
    for (unsigned int i = 0; i < 32; i++)
    {
        mbufs_to_send[i] = rte_pktmbuf_alloc(mbuf_pool);
        assert_non_null(mbufs_to_send[i]);
    }

    /* mbufs_to_send should be call as parameter for vnode_mirror_enqueue_bulk */
    expect_memory(__wrap_vnode_mirror_enqueue_bulk, buffs, mbufs_to_send, 32);
    expect_value(__wrap_vnode_mirror_enqueue_bulk, nr_buffs, 32);

    /* fault inject, all packet send successfully */
    will_return(__wrap_vnode_mirror_enqueue_bulk, 32);
    will_return(__wrap_vnode_mirror_rt_object_retrieve, NULL);
    will_return(__wrap_vnode_mirror_rt_object_retrieve, 0);

    /* call send with option, no options */
    int ret = marsio_send_burst_with_options(test_state->mock_sendpath, 0, (marsio_buff_t **)mbufs_to_send, 32, 0);
    assert_int_equal(ret, 0);

    /* make sure no mbuf leak */
    assert_int_equal(rte_mempool_in_use_count(mbuf_pool), 0);

    /* free the mbuf */
    for (unsigned int i = 0; i < 32; i++)
    {
        mbufs_to_send[i] = rte_pktmbuf_alloc(mbuf_pool);
        assert_non_null(mbufs_to_send[i]);
    }

    /* mbufs_to_send should be call as parameter for vnode_mirror_enqueue_bulk */
    expect_memory(__wrap_vnode_mirror_enqueue_bulk, buffs, mbufs_to_send, 32);
    expect_value(__wrap_vnode_mirror_enqueue_bulk, nr_buffs, 32);

    /* fault inject, all packet send successfully */
    will_return(__wrap_vnode_mirror_enqueue_bulk, 32);
    will_return(__wrap_vnode_mirror_rt_object_retrieve, NULL);
    will_return(__wrap_vnode_mirror_rt_object_retrieve, 0);

    /* first 16 mbufs, send by send burst */
    ret = marsio_send_burst(test_state->mock_sendpath, 0, (marsio_buff_t **)mbufs_to_send, 16);
    assert_int_equal(ret, 0);

    /* second 16 mbufs, send by send burst with options */
    ret = marsio_send_burst_with_options(test_state->mock_sendpath, 0, (marsio_buff_t **)mbufs_to_send + 16, 16, 0);
    assert_int_equal(ret, 0);

    /* poll and wait */
    struct mr_vdev * vdevs_handles[] = {test_state->mock_mr_vdev};
    marsio_poll_wait(test_state->mock_mr_vdev->instance, vdevs_handles, RTE_DIM(vdevs_handles), 0, 0);

    /* make sure no mbuf leak */
    assert_int_equal(rte_mempool_in_use_count(mbuf_pool), 0);
}
#endif

int main(int argc, char * argv[])
{
    /* generate the eal args */
    static char * eal_args[] = {
        "test_marsio_rawio_api",
        "-c 0x1",
        "--no-huge",
        "--no-pci",
    };

    /* run the eal with no huge pages */
    int ret = rte_eal_init(RTE_DIM(eal_args), eal_args);
    if (ret < 0)
    {
        rte_exit(EXIT_FAILURE, "Error with EAL initialization\n");
    }

    /* init the mbuf pool */
    mbuf_pool = rte_pktmbuf_pool_create("MBUF_POOL", 8192, 32, 0, RTE_MBUF_DEFAULT_BUF_SIZE, SOCKET_ID_ANY);

    const struct CMUnitTest tests[] = {
        cmocka_unit_test(test_recv_burst_api_one_pkt),
        cmocka_unit_test(test_recv_burst_api_many_pkts),
        cmocka_unit_test(test_send_burst_api_one_pkt_succ),
        cmocka_unit_test(test_send_burst_api_one_pkt_drop_by_backpressure),
    };

    const struct CMUnitTest testcase_send_with_options[] = {
        cmocka_unit_test_setup_teardown(test_send_with_options_do_not_free, NULL, NULL),
        cmocka_unit_test_setup_teardown(test_send_with_options_rehash, NULL, NULL),
        cmocka_unit_test_setup_teardown(test_send_with_options_many_pkts, NULL, NULL),
#if 0
        cmocka_unit_test_setup_teardown(test_mix_send_with_option_and_send_burst, NULL, NULL),
#endif
    };

    cmocka_run_group_tests(tests, NULL, NULL);
    cmocka_run_group_tests(testcase_send_with_options, testcase_group_SEND_WITH_OPTIONS_setup,
                           testcase_group_SEND_WITH_OPTIONS_teardown);
}