#pragma once

#include <rte_mbuf_core.h>
#include <rte_mbuf.h>
#include <rte_udp.h>
#include <rte_byteorder.h>
#include <rte_ip.h>
#include <stdint.h>
#include <rte_flow.h>
#include <rte_ether.h>
#include <unistd.h>
#include <rte_net.h>
#include <rte_vxlan.h>

constexpr static uint32_t ETHER_FRAME_MAGIC = 0xDCDCE5E5;
const static struct rte_ether_addr PROBE_MAC_ADDR {0x01,0x1B,0x19,0x00,0x00,0x00};
const static uint16_t ETHER_TYPE_LOCAL_EXP = 0x88b5;

struct ptp_hdr {
    uint8_t ptp_msg_type;
    uint8_t ptp_ver;
    uint8_t unused[34];
} __attribute__((packed));

struct pkt_hdr {
    struct rte_ether_hdr eth_hdr;
    struct ptp_hdr ptp_hdr;
    uint16_t type;
    uint32_t magic;
    char payload[0];
} __attribute__((packed));

constexpr static uint16_t PKT_TYPE_LOAD = 0;
constexpr static uint16_t PKT_TYPE_PROBE = 1;
constexpr static uint16_t PKT_TYPE_LOAD_RESP = 2;
constexpr static uint16_t PKT_TYPE_PROBE_RESP = 3;
struct pkt_payload_epoch {
    uint32_t epoch;
};

constexpr static uint16_t PKT_TYPE_STAT = 4;
struct pkt_payload_stat {
    uint32_t epoch;
    uint64_t hw_rx;
    uint64_t hw_tx;
    uint64_t sw_rx;
    uint64_t sw_tx;
};

constexpr static uint16_t NUM_PKT_TYPES = PKT_TYPE_STAT + 1;
// for fast packet verification
static const uint32_t expected_payload_size[NUM_PKT_TYPES] {
    sizeof(struct pkt_payload_epoch), // LOAD
    sizeof(struct pkt_payload_epoch), // PROBE
    sizeof(struct pkt_payload_epoch), // LOAD_RESP
    sizeof(struct pkt_payload_epoch), // PROBE_RESP
    sizeof(struct pkt_payload_stat) //STAT
};

static inline void
print_mac(struct rte_ether_addr * mac)
{
	printf("%x:%x:%x:%x:%x:%x", mac->addr_bytes[0],
								mac->addr_bytes[1],
								mac->addr_bytes[2],
								mac->addr_bytes[3],
								mac->addr_bytes[4],
								mac->addr_bytes[5]);
}

static inline void
print_ipv4(uint32_t ip)
{
	printf("%d-%d-%d-%d", (ip >> 24) & 0xff,
						  (ip >> 16) & 0xff,
						  (ip >> 8) & 0xff,
						  (ip >> 0) & 0xff);
}

static inline void
dump_pkt(struct rte_mbuf *pkt)
{
	if(rte_pktmbuf_data_len(pkt) < sizeof(struct rte_ether_hdr)) {
		return;
	}

	struct rte_ether_hdr _eth_hdr;
	struct rte_ether_hdr * eth_hdr = (struct rte_ether_hdr *)rte_pktmbuf_read(pkt, 0, sizeof(struct rte_ether_hdr), &_eth_hdr);
	if (eth_hdr == NULL) {
		return;
	}

	// ethernet frame
	printf("Packet %p: Length 0x%x\n", (void*)pkt, rte_pktmbuf_data_len(pkt));
	printf("    Ethernet header:\n");
	printf("        Src:");
	print_mac(&eth_hdr->s_addr);
	printf("\n");
	printf("        Dst:");
	print_mac(&eth_hdr->d_addr);
	printf("\n");
	printf("        Type: 0x%x\n", rte_be_to_cpu_16(eth_hdr->ether_type));

	uint16_t ether_type = rte_be_to_cpu_16(eth_hdr->ether_type);
	if (ether_type != RTE_ETHER_TYPE_IPV4) {
		return;
	}

	if(rte_pktmbuf_data_len(pkt) < sizeof(struct rte_ether_hdr) + sizeof(struct rte_ipv4_hdr)) {
		return;
	}

	// dump ip header
	struct rte_ipv4_hdr * ipv4_hdr = (struct rte_ipv4_hdr *)(eth_hdr + 1);	
	printf("    IPv4 header:\n");
	printf("        Src:");
	print_ipv4(rte_be_to_cpu_32(ipv4_hdr->src_addr));
	printf("\n");
	printf("        Dst:");
	print_ipv4(rte_be_to_cpu_32(ipv4_hdr->dst_addr));
	printf("\n");
	printf("        Protocol: 0x%x\n", ipv4_hdr->next_proto_id);

}


// fills the packet with the information except for the payload itself
static inline
struct pkt_hdr * construct_pkt_hdr(struct rte_mbuf * buf, uint16_t type,
                    struct rte_ether_addr * src_mac, struct rte_ether_addr * dst_mac)
{
    rte_pktmbuf_reset(buf);

    const uint32_t total_sz = sizeof(struct pkt_hdr) + expected_payload_size[type];
    struct pkt_hdr * pkt_data = (struct pkt_hdr *)rte_pktmbuf_append(buf, total_sz);
    struct rte_ether_hdr * eth_hdr;

    if (pkt_data == NULL)
        return NULL;

    // single segment
    buf->nb_segs = 1;
    
    // construct l2 header
    eth_hdr = &pkt_data->eth_hdr;
    rte_ether_addr_copy(src_mac, &eth_hdr->s_addr);
    if (type == PKT_TYPE_PROBE || type == PKT_TYPE_PROBE_RESP) {
        rte_ether_addr_copy(&PROBE_MAC_ADDR, &eth_hdr->d_addr);
        eth_hdr->ether_type = rte_cpu_to_be_16(RTE_ETHER_TYPE_1588);
        pkt_data->ptp_hdr.ptp_ver = 0x2; // VER 2
        buf->ol_flags |= PKT_TX_IEEE1588_TMST;
    } else {
        rte_ether_addr_copy(dst_mac, &eth_hdr->d_addr);
        eth_hdr->ether_type = rte_cpu_to_be_16(ETHER_TYPE_LOCAL_EXP);
        pkt_data->ptp_hdr.ptp_ver = 0xff;
    }
    buf->l2_len = sizeof(struct rte_ether_hdr);

    pkt_data->ptp_hdr.ptp_msg_type = 0x0; // SYNC
    pkt_data->type = rte_cpu_to_be_16(type);
    pkt_data->magic = rte_cpu_to_be_32(ETHER_FRAME_MAGIC);

    return pkt_data;
}

static inline
struct pkt_hdr * check_valid_packet(struct rte_mbuf * pkt)
{
    struct pkt_hdr * pkt_data = NULL;
    const uint32_t data_len = rte_pktmbuf_data_len(pkt);

    if (data_len < sizeof(struct pkt_hdr)) {
        return NULL;
    }

    pkt_data = rte_pktmbuf_mtod(pkt, struct pkt_hdr *);

    // check MAGIC
    if (rte_be_to_cpu_32(pkt_data->magic) != ETHER_FRAME_MAGIC) {
        return NULL;
    }

    // check type and payload size
    if ((rte_be_to_cpu_16(pkt_data->type) < NUM_PKT_TYPES) && 
        (data_len >= (sizeof(struct pkt_hdr) + expected_payload_size[rte_be_to_cpu_16(pkt_data->type)]))) {
        return pkt_data;
    }

    return NULL;
}