extern "C" {
#include <assert.h>
#include <marsio.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <rte_ether.h>
#include <rte_ip.h>
#include <rte_udp.h>
#include <rte_version.h>
}

#include <algorithm>

static char appsym[64] = "smartoffload";
static char dev_symbol[64] = "vxlan_user";

uint64_t cpu_mask = 0x1;
unsigned int nr_thread = 1;

struct mr_instance * mr_instance = NULL;
struct mr_vdev * dev_handler = NULL;
struct mr_sendpath * sendpath = NULL;

#define BURST_MAX 64
unsigned int nr_burst = 1;
unsigned int self_loop = 0;

int __is_bfd_pkt_and_make_answer(char * pkt_ptr, unsigned int pkt_len, bool make_answer)
{
    struct rte_ether_hdr * eth_hdr = (struct rte_ether_hdr *)pkt_ptr;

    /* must be ipv4 pkt for vxlan overlay */
    if (eth_hdr->ether_type != ntohs(RTE_ETHER_TYPE_IPV4))
    {
        return 0;
    }

    pkt_ptr += sizeof(struct rte_ether_hdr);
    pkt_len -= sizeof(struct rte_ether_hdr);

    /* ipv4 header */
    struct rte_ipv4_hdr * ipv4_hdr = (struct rte_ipv4_hdr *)pkt_ptr;
    if (ipv4_hdr->next_proto_id != IPPROTO_UDP)
    {
        return 0;
    }

    pkt_ptr += sizeof(struct rte_ipv4_hdr);
    pkt_len -= sizeof(struct rte_ipv4_hdr);

    struct rte_udp_hdr * udp_hdr = (struct rte_udp_hdr *)pkt_ptr;
    if (udp_hdr->dst_port != ntohs(3784))
    {
        return 0;
    }

    if (make_answer)
    {

        /* at here, this pkt must be a vxlan pkt,
           we need to swap the mac addresses and ipv4 addresses */

        struct rte_ether_addr swap_tmp;

#if RTE_VERSION_NUM(21, 11, 0, 0) <= RTE_VERSION
        rte_ether_addr_copy(&eth_hdr->dst_addr, &swap_tmp);
        rte_ether_addr_copy(&eth_hdr->src_addr, &eth_hdr->dst_addr);
        rte_ether_addr_copy(&swap_tmp, &eth_hdr->src_addr);
#else
        rte_ether_addr_copy(&eth_hdr->d_addr, &swap_tmp);
        rte_ether_addr_copy(&eth_hdr->s_addr, &eth_hdr->d_addr);
        rte_ether_addr_copy(&swap_tmp, &eth_hdr->s_addr);
#endif

        /* ipv4 src, dst and udp's ports must be swap */
        rte_be32_t ipv4_addr_swap_tmp;
        ipv4_addr_swap_tmp = ipv4_hdr->dst_addr;
        ipv4_hdr->dst_addr = ipv4_hdr->src_addr;
        ipv4_hdr->src_addr = ipv4_addr_swap_tmp;

        rte_be16_t udp_port_swap_tmp;
        udp_port_swap_tmp = udp_hdr->dst_port;
        udp_hdr->dst_port = udp_hdr->src_port;
        udp_hdr->src_port = udp_port_swap_tmp;
    }

    return 1;
}

int __is_g_vxlan_and_make_inject(char * pkt_ptr, unsigned int pkt_len, bool do_inject)
{
    struct rte_ether_hdr * eth_hdr = (struct rte_ether_hdr *)pkt_ptr;

    /* must be ipv4 pkt for vxlan overlay */
    if (eth_hdr->ether_type != ntohs(RTE_ETHER_TYPE_IPV4))
    {
        return 0;
    }

    pkt_ptr += sizeof(struct rte_ether_hdr);
    pkt_len -= sizeof(struct rte_ether_hdr);

    /* ipv4 header */
    struct rte_ipv4_hdr * ipv4_hdr = (struct rte_ipv4_hdr *)pkt_ptr;
    if (ipv4_hdr->next_proto_id != IPPROTO_UDP)
    {
        return 0;
    }

    pkt_ptr += sizeof(struct rte_ipv4_hdr);
    pkt_len -= sizeof(struct rte_ipv4_hdr);

    struct rte_udp_hdr * udp_hdr = (struct rte_udp_hdr *)pkt_ptr;
    if (udp_hdr->dst_port != ntohs(4789))
    {
        return 0;
    }

    /* at here, this pkt must be a vxlan pkt,
       we need to swap the mac addresses and ipv4 addresses */

    if (do_inject)
    {
        struct rte_ether_addr swap_tmp;

#if RTE_VERSION_NUM(21, 11, 0, 0) <= RTE_VERSION
        rte_ether_addr_copy(&eth_hdr->dst_addr, &swap_tmp);
        rte_ether_addr_copy(&eth_hdr->src_addr, &eth_hdr->dst_addr);
        rte_ether_addr_copy(&swap_tmp, &eth_hdr->src_addr);
#else
        rte_ether_addr_copy(&eth_hdr->d_addr, &swap_tmp);
        rte_ether_addr_copy(&eth_hdr->s_addr, &eth_hdr->d_addr);
        rte_ether_addr_copy(&swap_tmp, &eth_hdr->s_addr);
#endif

        rte_be32_t ipv4_addr_swap_tmp;
        ipv4_addr_swap_tmp = ipv4_hdr->dst_addr;
        ipv4_hdr->dst_addr = ipv4_hdr->src_addr;
        ipv4_hdr->src_addr = ipv4_addr_swap_tmp;
    }

    return 1;
}

void * txonly_loop(void * arg)
{
    uintptr_t sid = (uintptr_t)arg;
    marsio_buff_t * rx_buff[BURST_MAX];
    marsio_thread_init(mr_instance);

    for (;;)
    {
        int ret = marsio_recv_burst(dev_handler, sid, rx_buff, 1);
        if (ret != 1)
            continue;

        marsio_buff_t * rx_buff_ptr = rx_buff[0];
        char * pkt_ptr = marsio_buff_mtod(rx_buff_ptr);
        unsigned int pkt_len = marsio_buff_datalen(rx_buff_ptr);
        assert(pkt_ptr != NULL && pkt_len != 0);

        bool should_inject = false;

        if (__is_bfd_pkt_and_make_answer(pkt_ptr, pkt_len, 0) > 0)
        {
            __is_bfd_pkt_and_make_answer(pkt_ptr, pkt_len, 1);
            should_inject = true;
        }
        else if (__is_g_vxlan_and_make_inject(pkt_ptr, pkt_len, 0) > 0)
        {
            /* for all vxlan pkts, create an offload request */
            marsio_buff_t * offload_request_buf = marsio_buff_malloc_smartoffload(dev_handler, pkt_ptr, pkt_len);
            assert(offload_request_buf != NULL);

            __is_g_vxlan_and_make_inject(pkt_ptr, pkt_len, 1);

            /* inject the offload request */
            marsio_send_burst_with_options(sendpath, sid, &offload_request_buf, 1, MARSIO_SEND_OPT_CTRL);
            should_inject = true;
        }

        if (should_inject)
        {
            marsio_send_burst_with_options(sendpath, sid, &rx_buff_ptr, 1, 0);
        }
        else
        {
            marsio_buff_free(mr_instance, rx_buff, 1, MARSIO_SOCKET_ID_ANY, MARSIO_LCORE_ID_ANY);
        }
    }

    return (void *)NULL;
}

int help()
{
    return 0;
}

int main(int argc, char * argv[])
{
    int opt = 0;
    while ((opt = getopt(argc, argv, "s:t:a:c:b:d:h?rl")) != -1)
    {
        char * endptr = NULL;
        switch (opt)
        {
        case '?':
        case 'h': {
            help();
            break;
        }
        case 'd': {
            snprintf(dev_symbol, sizeof(dev_symbol), "%s", optarg);
            break;
        }
        case 'a': {
            snprintf(appsym, sizeof(appsym), "%s", optarg);
            break;
        }
        case 'c': {
            cpu_mask = strtoull(optarg, &endptr, 0);
            if (cpu_mask == 0 && endptr == optarg)
                help();
            break;
        }
        case 'b': {
            nr_burst = strtoull(optarg, &endptr, 0);
            if (nr_burst == 0 && endptr == optarg)
                help();
            break;
        }
        case 'l': {
            self_loop = 1;
            break;
        }

        default:
            help();
            break;
        }
    }

    mr_instance = marsio_create();
    if (mr_instance == NULL)
    {
        fprintf(stderr, "Marsio instance create failed. ");
        abort();
    }

    unsigned int opt_value = 1;
    marsio_option_set(mr_instance, MARSIO_OPT_EXIT_WHEN_ERR, &opt_value, sizeof(opt_value));
    marsio_option_set(mr_instance, MARSIO_OPT_THREAD_MASK, &cpu_mask, sizeof(cpu_mask));
    marsio_init(mr_instance, appsym);

    nr_thread = __builtin_popcountll(cpu_mask);
    dev_handler = marsio_open_device(mr_instance, dev_symbol, nr_thread, nr_thread);
    fprintf(stdout, "Thread Count = %d\n", nr_thread);

    sendpath = marsio_sendpath_create_by_vdev(dev_handler);
    assert(sendpath != NULL);

    pthread_t __tmp_pid[nr_thread];
    for (int i = 0; i < nr_thread; i++)
    {
        pthread_create(&__tmp_pid[i], NULL, txonly_loop, (void *)(uintptr_t)i);
    }

    for (int i = 0; i < nr_thread; i++)
    {
        pthread_join(__tmp_pid[i], NULL);
    }

    marsio_destory(mr_instance);
    fprintf(stdout, "RXONLY is terminated. ");
    return 0;
}