#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <inttypes.h>
#include <sys/types.h>
#include <string.h>
#include <sys/queue.h>
#include <stdarg.h>
#include <errno.h>
#include <getopt.h>
#include <signal.h>
#include <sys/param.h>

#include <rte_common.h>
#include <rte_byteorder.h>
#include <rte_log.h>
#include <rte_memory.h>
#include <rte_memcpy.h>
#include <rte_eal.h>
#include <rte_launch.h>
#include <rte_atomic.h>
#include <rte_cycles.h>
#include <rte_prefetch.h>
#include <rte_lcore.h>
#include <rte_per_lcore.h>
#include <rte_branch_prediction.h>
#include <rte_interrupts.h>
#include <rte_random.h>
#include <rte_debug.h>
#include <rte_ether.h>
#include <rte_ethdev.h>
#include <rte_mempool.h>
#include <rte_mbuf.h>
#include <rte_malloc.h>
#include <rte_ip.h>
#include <rte_tcp.h>
#include <rte_udp.h>
#include <rte_string_fns.h>
#include <rte_lpm.h>
#include <rte_lpm6.h>

#include <rte_ip_frag.h>

#define DEF_FLOW_NUM    0x1000
#define DEF_FLOW_TTL    0x100000
#define IP_FRAG_TBL_BUCKET_ENTRIES  128

#define RX_RING_SIZE 1024
#define TX_RING_SIZE 1024

#define NUM_MBUFS 8191
#define MBUF_CACHE_SIZE 250
#define BURST_SIZE 32

static volatile uint8_t force_quit;

static const struct rte_eth_conf port_conf_default = {
        .rxmode = {
                .max_rx_pkt_len = RTE_ETHER_MAX_LEN,
        },
};


static uint32_t max_flow_num = DEF_FLOW_NUM;
static uint32_t max_flow_ttl = DEF_FLOW_TTL;

struct lcore_queue_conf {
    struct rte_ip_frag_tbl *frag_tbl;
    struct rte_mempool *pool;
    struct rte_ip_frag_death_row death_row; 
}__rte_cache_aligned; 
static struct lcore_queue_conf lcore_queue_conf[RTE_MAX_LCORE];


static void
signal_handler(int signum)
{
        if (signum == SIGINT || signum == SIGTERM) {
                printf("\n\nSignal %d received, preparing to exit...\n",
                                signum);
                force_quit = 1;
        }
}


static inline int setup_queue_tbl(struct lcore_queue_conf *qconf)
{

    uint64_t frag_cycles = (rte_get_tsc_hz() + MS_PER_S - 1) / MS_PER_S * max_flow_ttl;
    
    qconf->frag_tbl = rte_ip_frag_table_create(max_flow_num, IP_FRAG_TBL_BUCKET_ENTRIES, max_flow_num, frag_cycles, rte_socket_id());

    if((qconf->frag_tbl) == NULL){ 
        printf("ERR, IP_RSMBL, Table Failed.");
        return -1;
    }   

    qconf->pool=rte_pktmbuf_pool_create("BUFFER", 8000, 256, 0, RTE_MBUF_DEFAULT_BUF_SIZE, rte_socket_id());

    if(qconf->pool== NULL){ 
	printf("ERR, IP_RSMBL, Mem Pool Failed.");
	return -1;
    }

    return 0;
}


static inline void reassemble(struct rte_mbuf *reassemblepkt, struct lcore_queue_conf *qconf, uint64_t cur_tsc)
{
    struct rte_mbuf *mo;
    struct rte_ether_hdr *eth_hdr;
    struct rte_ipv4_hdr *ip_hdr;

    eth_hdr = rte_pktmbuf_mtod(reassemblepkt, struct rte_ether_hdr *);
    ip_hdr = (struct rte_ipv4_hdr *)(eth_hdr + 1);

    if (rte_ipv4_frag_pkt_is_fragmented(ip_hdr)){
        reassemblepkt->l2_len = sizeof(*eth_hdr);
        reassemblepkt->l3_len = sizeof(*ip_hdr);
        
        int ip_len;
        ip_len = rte_be_to_cpu_16(ip_hdr->total_length);

        mo = rte_ipv4_frag_reassemble_packet(qconf->frag_tbl, &qconf->death_row, reassemblepkt, cur_tsc, ip_hdr);
            
        if (mo == NULL){
	        printf("Total Length: %u, l3 length: %u, Packet ID: %u \n", ip_len, reassemblepkt->l3_len, ip_hdr->packet_id);
		printf("src add: %x dst add: %x \n", ip_hdr->src_addr,ip_hdr->dst_addr); 
        	printf("ERR, IP_RSMBL, Reassemble Failed.\n");
        }   

        if ((mo != reassemblepkt) && (mo != NULL)){
            printf("Reassemble is success.");
            reassemblepkt = mo;

	rte_pktmbuf_dump(stdout, mo, 64);
	rte_pktmbuf_free(mo);
        }
    }
}

static int
lcore_main(struct lcore_queue_conf *qconf)
{
#define BUFFER_LENGTH 32
    int rx, rec;
    struct rte_mbuf *bufs[BUFFER_LENGTH];
    uint64_t cur_tsc;   
    int i;

	uint16_t port;

    RTE_ETH_FOREACH_DEV(port){
    printf("RX Thread: Socket ID: %u\n", rte_socket_id());
    printf("RX Thread: lcore count: %u\n", rte_lcore_count());
    printf("RX Thread: lcore ID: %u\n", rte_lcore_id()); 
    }   
        
        while(!force_quit) {
            cur_tsc = rte_rdtsc();
            RTE_ETH_FOREACH_DEV(port){
                            rx=rte_eth_rx_burst(port, 0, bufs, BUFFER_LENGTH); 
                            if(likely(rx)) { 
                                for(i=0; i<rx; i++)
	                                reassemble(bufs[i], qconf, cur_tsc);

                                rte_ip_frag_free_death_row(&qconf->death_row, 2);
                            }
            }
        }
        return 0;
}

int port_init(uint16_t port, struct rte_mempool *mbuf_pool)
{
        struct rte_eth_conf port_conf = port_conf_default;
        const uint16_t rx_rings = 1, tx_rings = 1;
        uint16_t nb_rxd = RX_RING_SIZE;
        uint16_t nb_txd = TX_RING_SIZE;
        int retval;
        uint16_t q;
        struct rte_eth_dev_info dev_info;
        struct rte_eth_txconf txconf;

        if (!rte_eth_dev_is_valid_port(port))
                return -1;

        retval = rte_eth_dev_info_get(port, &dev_info);
        if (retval != 0) {
                printf("Error during getting device (port %u) info: %s\n",
                                port, strerror(-retval));
                return retval;
        }

        if (dev_info.tx_offload_capa & DEV_TX_OFFLOAD_MBUF_FAST_FREE)
                port_conf.txmode.offloads |=
                        DEV_TX_OFFLOAD_MBUF_FAST_FREE;

        /* Configure the Ethernet device. */
        retval = rte_eth_dev_configure(port, rx_rings, tx_rings, &port_conf);
        if (retval != 0)
                return retval;

        retval = rte_eth_dev_adjust_nb_rx_tx_desc(port, &nb_rxd, &nb_txd);
        if (retval != 0)
                return retval;

        /* Allocate and set up 1 RX queue per Ethernet port. */
        for (q = 0; q < rx_rings; q++) {
                retval = rte_eth_rx_queue_setup(port, q, nb_rxd,
                                rte_eth_dev_socket_id(port), NULL, mbuf_pool);
                if (retval < 0)
                        return retval;
        }
        txconf = dev_info.default_txconf;
        txconf.offloads = port_conf.txmode.offloads;
        /* Allocate and set up 1 TX queue per Ethernet port. */
        for (q = 0; q < tx_rings; q++) {
                retval = rte_eth_tx_queue_setup(port, q, nb_txd,
                                rte_eth_dev_socket_id(port), &txconf);
                if (retval < 0)
                        return retval;
        }

        /* Start the Ethernet port. */
        retval = rte_eth_dev_start(port);
        if (retval < 0)
                return retval;

        /* Display the port MAC address. */
        struct rte_ether_addr addr;
        retval = rte_eth_macaddr_get(port, &addr);
        if (retval != 0)
                return retval;

        printf("Port %u MAC: %02" PRIx8 " %02" PRIx8 " %02" PRIx8
                           " %02" PRIx8 " %02" PRIx8 " %02" PRIx8 "\n",
                        port,
                        addr.addr_bytes[0], addr.addr_bytes[1],
                        addr.addr_bytes[2], addr.addr_bytes[3],
                        addr.addr_bytes[4], addr.addr_bytes[5]);

        /* Enable RX in promiscuous mode for the Ethernet device. */
        retval = rte_eth_promiscuous_enable(port);
        if (retval != 0)
                return retval;

        return 0;
}


int main(int argc, char *argv[])
{

    int ret;
    uint16_t portcheck;
    uint16_t p1;
    struct lcore_queue_conf *qconf;

    force_quit = 0;
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);

    /* EAL setup */
    ret=rte_eal_init(argc, argv);
    if(ret < 0)
        rte_exit(EXIT_FAILURE, "EAL initialising failed.");

    qconf = &lcore_queue_conf[rte_get_master_lcore()];

    if(setup_queue_tbl(qconf) != 0)
    rte_exit(EXIT_FAILURE, "%s\n", rte_strerror(rte_errno));

	RTE_ETH_FOREACH_DEV(portcheck){
	    if(port_init(portcheck, qconf->pool) != 0)
		rte_exit(EXIT_FAILURE, "Ethernet port initialisation failed.");
    }


    /* Master core call */  
    lcore_main(qconf);

    return 0;
}