#include <setjmp.h>
#include <stdarg.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>

#include <cmocka.h>

#include <rte_eal.h>
#include <rte_graph.h>
#include <rte_graph_worker.h>
#include <rte_mempool.h>

#include <sc_node.h>
#include <sc_node_common.h>

static struct rte_mempool * mock_pktmbuf_pool = NULL;

int __wrap_rte_eal_init(int argc, char * argv[])
{
    int inject_eal_init_failed = mock_type(int);
    if (inject_eal_init_failed)
    {
        return -1;
    }

    return 0;
}

static int rte_exit_called = 0;
int __wrap_rte_exit(int status, const char * format, ...)
{
    rte_exit_called = 1;
    return 0;
}

struct mr_dev_desc * __wrap_mr_dev_desc_lookup(struct devmgr_main * devmgr_main, const char * devsym)
{
    check_expected(devsym);
    return mock_type(struct mr_dev_desc *);
}

void __wrap_forwarder_table_insert(uint16_t sid, uint16_t type)
{
    check_expected(sid);
    check_expected(type);
}

extern int __real_rte_eal_init(int argc, char ** argv);
extern int __real_rte_exit(int status, const char * format, ...);

extern int vwire_init(struct sc_main * sc);

/* testcase: 1
 * device: vwire_1 <-> vwire_2
 * port_id: 1 <-> 2
 */

static void testcase_vwire_ingress_forward_pkts(void ** state)
{
    struct rte_graph * graph_ptr = rte_graph_lookup("test_vwire_ingress");
    assert_non_null(graph_ptr);

    /* alloc the mbuf as input */
    struct rte_mbuf * mbuf_input[3];
    mbuf_input[0] = rte_pktmbuf_alloc(mock_pktmbuf_pool);
    mbuf_input[1] = rte_pktmbuf_alloc(mock_pktmbuf_pool);
    mbuf_input[2] = rte_pktmbuf_alloc(mock_pktmbuf_pool);

    assert_non_null(mbuf_input[0]);
    assert_non_null(mbuf_input[1]);
    assert_non_null(mbuf_input[2]);

    /* set the mbuf data */
    will_return(mock_pkt_input_process, 3);
    will_return(mock_pkt_input_process, mbuf_input);

    struct rte_mbuf * mbuf_output[2] = {};
    unsigned int nr_mbuf_output = 0;

    will_return(mock_pkt_output_process, mbuf_output);
    will_return(mock_pkt_output_process, &nr_mbuf_output);

    struct rte_mbuf * mbuf_drop[1] = {};
    unsigned int nr_mbuf_drop = 0;
    will_return(mock_pkt_drop_process, mbuf_drop);
    will_return(mock_pkt_drop_process, &nr_mbuf_drop);

    /* first pkt, I->E */
    struct mrb_metadata * mrb_meta_0 = mrbuf_cz_data(mbuf_input[0], MR_NODE_CTRLZONE_ID);
    mrb_metadata_clear(mrb_meta_0);
    sid_list_init(&mrb_meta_0->sid_list);
    mrb_meta_0->port_ingress = 0;

    /* second pkt, E->I */
    struct mrb_metadata * mrb_meta_1 = mrbuf_cz_data(mbuf_input[1], MR_NODE_CTRLZONE_ID);
    mrb_metadata_clear(mrb_meta_1);
    sid_list_init(&mrb_meta_1->sid_list);
    mrb_meta_1->port_ingress = 1;

    /* thrid pkt, from non vwire device, should be drop */
    struct mrb_metadata * mrb_meta_2 = mrbuf_cz_data(mbuf_input[2], MR_NODE_CTRLZONE_ID);
    mrb_metadata_clear(mrb_meta_2);
    sid_list_init(&mrb_meta_2->sid_list);
    mrb_meta_2->port_ingress = 2;

    rte_graph_walk(graph_ptr);

    /* check the count */
    assert_int_equal(nr_mbuf_output, 2);
    assert_int_equal(nr_mbuf_drop, 1);

    /* check the mbuf ptr */
    assert_ptr_equal(mbuf_output[0], mbuf_input[0]);
    assert_ptr_equal(mbuf_output[1], mbuf_input[1]);
    assert_ptr_equal(mbuf_drop[0], mbuf_input[2]);

    /* check the dir */
    assert_int_equal(mrb_meta_0->dir, 0);
    assert_int_equal(mrb_meta_1->dir, 1);

    /* check the sid */
    uint16_t mbuf_0_sids[MR_NODE_COMMON_MAX_SID_NUM] = {};
    unsigned int nr_mbuf_0_sids = 0;

    nr_mbuf_0_sids = sid_list_get(&mrb_meta_0->sid_list, mbuf_0_sids, RTE_DIM(mbuf_0_sids));
    assert_int_equal(nr_mbuf_0_sids, 1);
    assert_int_equal(mbuf_0_sids[0], 300);

    /* check the second packet */
    uint16_t mbuf_1_sids[MR_NODE_COMMON_MAX_SID_NUM] = {};
    unsigned int nr_mbuf_1_sids = 0;

    nr_mbuf_1_sids = sid_list_get(&mrb_meta_1->sid_list, mbuf_1_sids, RTE_DIM(mbuf_1_sids));
    assert_int_equal(nr_mbuf_1_sids, 1);
    assert_int_equal(mbuf_1_sids[0], 300);

    /* release the mbufs */
    rte_pktmbuf_free_bulk(mbuf_output, 2);
    rte_pktmbuf_free_bulk(mbuf_drop, 1);

    assert_true(rte_mempool_full(mock_pktmbuf_pool) == 1);
}

static void testcase_vwire_egress_forward_pkts(void ** state)
{
    struct rte_graph * graph_ptr = rte_graph_lookup("test_vwire_egress");
    assert_non_null(graph_ptr);

    /* alloc the mbuf as input */
    struct rte_mbuf * mbuf_input[2];
    mbuf_input[0] = rte_pktmbuf_alloc(mock_pktmbuf_pool);
    mbuf_input[1] = rte_pktmbuf_alloc(mock_pktmbuf_pool);

    assert_non_null(mbuf_input[0]);
    assert_non_null(mbuf_input[1]);

    /* set the mbuf data */
    will_return(mock_pkt_input_process, 2);
    will_return(mock_pkt_input_process, mbuf_input);

    struct rte_mbuf * mbuf_output[2] = {};
    unsigned int nr_mbuf_output = 0;

    will_return(mock_pkt_output_process, mbuf_output);
    will_return(mock_pkt_output_process, &nr_mbuf_output);

    uint16_t sid_list_to_prepend[1] = {300};

    /* first pkt, I->E */
    struct mrb_metadata * mrb_meta_0 = mrbuf_cz_data(mbuf_input[0], MR_NODE_CTRLZONE_ID);
    mrb_metadata_clear(mrb_meta_0);
    sid_list_init(&mrb_meta_0->sid_list);
    sid_list_prepend(&mrb_meta_0->sid_list, sid_list_to_prepend, RTE_DIM(sid_list_to_prepend));
    mrb_meta_0->cur_sid = 300;
    mrb_meta_0->dir = 0;

    /* second pkt, E->I */
    struct mrb_metadata * mrb_meta_1 = mrbuf_cz_data(mbuf_input[1], MR_NODE_CTRLZONE_ID);
    mrb_metadata_clear(mrb_meta_1);
    sid_list_init(&mrb_meta_1->sid_list);
    mrb_meta_1->cur_sid = 300;
    mrb_meta_1->dir = 1;

    rte_graph_walk(graph_ptr);

    /* check the count */
    assert_int_equal(nr_mbuf_output, 2);

    /* check the mbuf ptr */
    assert_ptr_equal(mbuf_output[0], mbuf_input[0]);
    assert_ptr_equal(mbuf_output[1], mbuf_input[1]);

    /* check the dir */
    assert_int_equal(mrb_meta_0->dir, 0);
    assert_int_equal(mrb_meta_1->dir, 1);
    assert_int_equal(mrb_meta_0->port_egress, 1);
    assert_int_equal(mrb_meta_1->port_egress, 0);

    /* release the buffs */
    rte_pktmbuf_free_bulk(mbuf_output, 2);
    assert_true(rte_mempool_full(mock_pktmbuf_pool) == 1);
}

static int testgroup_vwire_egress_setup(void ** state)
{
    /* mock pkt source to generate mock packets */
    const char * node_patterns[] = {
        "mock_pkt_input",
        "mock_pkt_output",
        "vwire_egress",
    };

    /* create a fake graph */
    struct rte_graph_param graph_param = {
        .node_patterns = node_patterns,
        .nb_node_patterns = RTE_DIM(node_patterns),
        .socket_id = 0,
    };

    /* update the pkt input's next id */
    rte_node_t node_id_pkt_input = rte_node_from_name("mock_pkt_input");
    assert_int_not_equal(node_id_pkt_input, RTE_NODE_ID_INVALID);

    const char * updated_next_node[] = {"vwire_egress"};
    int rc = rte_node_edge_update(node_id_pkt_input, 0, updated_next_node, RTE_DIM(updated_next_node));
    assert_int_equal(rc, 1);

    /* update the node edge */
    rte_node_t node_id_vwire_egress = rte_node_from_name("vwire_egress");
    assert_int_not_equal(node_id_vwire_egress, RTE_NODE_ID_INVALID);

    const char * vwire_egress_updated_next_node[] = {
        "mock_pkt_output",
        "mock_pkt_drop",
    };

    rc = rte_node_edge_update(node_id_vwire_egress, 0, vwire_egress_updated_next_node,
                              RTE_DIM(vwire_egress_updated_next_node));

    assert_int_equal(rc, 2);

    rte_graph_t graph = rte_graph_create("test_vwire_egress", &graph_param);
    assert_int_not_equal(graph, RTE_GRAPH_ID_INVALID);

    *state = (void *)(uintptr_t)graph;
    return 0;
}

static int testgroup_vwire_egress_teardown(void ** state)
{
    /* destroy the graph */
    rte_graph_t graph = (uintptr_t)*state;
    return rte_graph_destroy(graph);
}

static int testgroup_vwire_ingress_setup(void ** state)
{
    /* mock pkt source to generate mock packets */
    const char * node_patterns[] = {
        "mock_pkt_input",
        "mock_pkt_output",
        "vwire_ingress",
    };

    /* create a fake graph */
    struct rte_graph_param graph_param = {
        .node_patterns = node_patterns,
        .nb_node_patterns = RTE_DIM(node_patterns),
        .socket_id = 0,
    };

    /* update the node edge */
    rte_node_t node_id = rte_node_from_name("vwire_ingress");
    assert_int_not_equal(node_id, RTE_NODE_ID_INVALID);

    rte_node_t next_node_id = rte_node_from_name("mock_pkt_output");
    assert_int_not_equal(next_node_id, RTE_NODE_ID_INVALID);

    const char * updated_next_node[] = {
        "mock_pkt_output",
        "mock_pkt_drop",
    };

    int rc = rte_node_edge_update(node_id, 0, updated_next_node, RTE_DIM(updated_next_node));
    assert_int_equal(rc, 2);

    rte_graph_t graph = rte_graph_create("test_vwire_ingress", &graph_param);
    assert_int_not_equal(graph, RTE_GRAPH_ID_INVALID);

    *state = (void *)(uintptr_t)graph;
    return 0;
}

static int testgroup_vwire_ingress_teardown(void ** state)
{
    /* destroy the graph */
    rte_graph_t graph = (uintptr_t)*state;
    return rte_graph_destroy(graph);
}

/* ----------------------------------------------------------------------------------------------------------------- */
static void testcase_vwire_init_success(void ** state)
{
    struct sc_main * mock_sc = *state;

    struct mr_dev_desc * mock_dev_desc_1 = rte_zmalloc(NULL, sizeof(struct mr_dev_desc), 0);
    assert_non_null(mock_dev_desc_1);

    struct mr_dev_desc * mock_dev_desc_2 = rte_zmalloc(NULL, sizeof(struct mr_dev_desc), 0);
    assert_non_null(mock_dev_desc_2);

    mock_dev_desc_1->port_id = 0;
    mock_dev_desc_2->port_id = 1;

    /* find vwire_1 device */
    expect_string(__wrap_mr_dev_desc_lookup, devsym, "vwire_1");
    will_return(__wrap_mr_dev_desc_lookup, mock_dev_desc_1);

    /* find vwire_2 device */
    expect_string(__wrap_mr_dev_desc_lookup, devsym, "vwire_2");
    will_return(__wrap_mr_dev_desc_lookup, mock_dev_desc_2);

    /* expect sid register */
    expect_value(__wrap_forwarder_table_insert, sid, 300);
    expect_value(__wrap_forwarder_table_insert, type, FORWARDER_TYPE_VWIRE);

    /* init should success */
    int ret = vwire_init(mock_sc);
    assert_int_equal(ret, RT_SUCCESS);

    rte_free(mock_dev_desc_1);
    rte_free(mock_dev_desc_2);
}

static void testcase_vwire_device_not_exist(void ** state)
{
    struct sc_main * mock_sc = *state;

    /* cannot find vwire_1 device */
    expect_string(__wrap_mr_dev_desc_lookup, devsym, "vwire_1");
    will_return(__wrap_mr_dev_desc_lookup, NULL);

    /* init should fail */
    int ret = vwire_init(mock_sc);
    assert_int_equal(ret, RT_ERR);
}

static int testgroup_vwire_init_setup(void ** state)
{
    struct sc_main * mock_sc = rte_zmalloc(NULL, sizeof(struct sc_main), 0);
    assert_non_null(mock_sc);

    remove("test_node_vwire.ini");
    FILE * fp_cfgfile = fopen("test_node_vwire.ini", "w");
    assert_non_null(fp_cfgfile);

    const char * test_cfgfile_content = "[vwire:0]\n"
                                        "vwire_id=0\n"
                                        "interface_int=vwire_1\n"
                                        "interface_ext=vwire_2";

    fprintf(fp_cfgfile, "%s", test_cfgfile_content);
    fclose(fp_cfgfile);

    strcpy(mock_sc->local_cfgfile, "test_node_vwire.ini");
    *state = mock_sc;
    return 0;
}

static int testgroup_vwire_init_teardown(void ** state)
{
    struct sc_main * mock_sc = *state;
    rte_free(mock_sc);
    unlink(mock_sc->local_cfgfile);
    return 0;
}

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

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

    mock_pktmbuf_pool = rte_pktmbuf_pool_create("test_pktmbuf_pool", 1024, 32, 128, RTE_MBUF_DEFAULT_BUF_SIZE,
                                                SOCKET_ID_ANY);

    const struct CMUnitTest group_vwire_init[] = {
        cmocka_unit_test(testcase_vwire_device_not_exist),
        cmocka_unit_test(testcase_vwire_init_success),
    };

    const struct CMUnitTest group_vwire_ingress[] = {
        cmocka_unit_test(testcase_vwire_ingress_forward_pkts),
    };

    const struct CMUnitTest group_vwire_egress[] = {
        cmocka_unit_test(testcase_vwire_egress_forward_pkts),
    };

    ret = cmocka_run_group_tests(group_vwire_init, testgroup_vwire_init_setup, testgroup_vwire_init_teardown);
    ret += cmocka_run_group_tests(group_vwire_ingress, testgroup_vwire_ingress_setup, testgroup_vwire_ingress_teardown);
    ret += cmocka_run_group_tests(group_vwire_egress, testgroup_vwire_egress_setup, testgroup_vwire_egress_teardown);
    rte_eal_cleanup();

    return ret;
}