#ifndef _PACKET_PUB_H
#define _PACKET_PUB_H

#ifdef __cplusplus
extern "C"
{
#endif

#include <string.h>
#include <arpa/inet.h>
#define __FAVOR_BSD 1
#include <netinet/tcp.h>
#include <netinet/udp.h>
#include <netinet/ip.h>
#include <netinet/ip6.h>

enum layer_type
{
    // L2 -- data link layer
    LAYER_TYPE_ETHER = 1 << 0,
    LAYER_TYPE_PWETH = 1 << 1,
    LAYER_TYPE_PPP = 1 << 2,
    LAYER_TYPE_HDLC = 1 << 3,
    LAYER_TYPE_L2TP = 1 << 4,
    LAYER_TYPE_L2 = (LAYER_TYPE_ETHER | LAYER_TYPE_PWETH | LAYER_TYPE_PPP | LAYER_TYPE_HDLC | LAYER_TYPE_L2TP),

    // L2 -- tunnel
    LAYER_TYPE_VLAN = 1 << 5,
    LAYER_TYPE_PPPOE = 1 << 6,
    LAYER_TYPE_MPLS = 1 << 7,
    LAYER_TYPE_L2_TUN = (LAYER_TYPE_VLAN | LAYER_TYPE_PPPOE | LAYER_TYPE_MPLS),

    // L3 -- network layer
    LAYER_TYPE_IPV4 = 1 << 8,
    LAYER_TYPE_IPV6 = 1 << 9,
    LAYER_TYPE_IPAH = 1 << 10,
    LAYER_TYPE_L3 = (LAYER_TYPE_IPV4 | LAYER_TYPE_IPV6 | LAYER_TYPE_IPAH),

    // L3 -- tunnel
    LAYER_TYPE_GRE = 1 << 11,
    LAYER_TYPE_L3_TUN = (LAYER_TYPE_GRE),

    // L4 -- transport layer
    LAYER_TYPE_UDP = 1 << 12,
    LAYER_TYPE_TCP = 1 << 13,
    LAYER_TYPE_ICMP = 1 << 14,
    LAYER_TYPE_ICMP6 = 1 << 15,
    LAYER_TYPE_L4 = (LAYER_TYPE_UDP | LAYER_TYPE_TCP | LAYER_TYPE_ICMP | LAYER_TYPE_ICMP6),

    // L4 -- tunnel
    LAYER_TYPE_VXLAN = 1 << 16,
    LAYER_TYPE_GTPV1_U = 1 << 17,

    // ALL
    LAYER_TYPE_ALL = (LAYER_TYPE_L2 | LAYER_TYPE_L2_TUN | LAYER_TYPE_L3 | LAYER_TYPE_L3_TUN | LAYER_TYPE_L4 | LAYER_TYPE_VXLAN | LAYER_TYPE_GTPV1_U),
};

struct packet_layer
{
    enum layer_type type;
    const char *hdr_ptr; // header pointer
    const char *pld_ptr; // payload pointer
    uint16_t hdr_offset; // header offset from data_ptr
    uint16_t hdr_len;    // header length
    uint16_t pld_len;    // payload length
};

#define MAX_SID_NUM 8
struct sid_list
{
    uint16_t sid[MAX_SID_NUM];
    int used;
};

enum packet_direction
{
    PACKET_DIRECTION_OUTGOING = 0, // Internal -> External: 0
    PACKET_DIRECTION_INCOMING = 1, // External -> Internal: 1
};

enum packet_action
{
    PACKET_ACTION_FORWARD = 0,
    PACKET_ACTION_DROP = 1,
};

enum packet_direction packet_get_direction(const struct packet *pkt);
uint64_t packet_get_session_id(const struct packet *pkt);
void packet_prepend_sid_list(struct packet *pkt, const struct sid_list *list);

int8_t packet_get_layers_number(const struct packet *pkt);
const struct packet_layer *packet_get_layer(const struct packet *pkt, int8_t idx);

const char *packet_get_data(const struct packet *pkt);
uint16_t packet_get_len(const struct packet *pkt);

const char *packet_get_payload(const struct packet *pkt);
uint16_t packet_get_payload_len(const struct packet *pkt);

void packet_set_action(struct packet *pkt, enum packet_action action);
enum packet_action packet_get_action(const struct packet *pkt);

/*
 ******************************************************************************
 * Example: foreach layer in packet
 ******************************************************************************
 *
 *    // inorder
 *    int8_t layers = packet_get_layers_number(pkt);
 *    for (int8_t i = 0; i < layers; i++)
 *    {
 *        const struct packet_layer *layer = packet_get_layer(pkt, i);
 *        printf("layer[%d]: type=%d, hdr_offset=%d, hdr_len=%d, pld_len=%d\n", i, layer->type, layer->hdr_offset, layer->hdr_len, layer->pld_len);
 *    }
 *
 *    // reverse
 *    for (int8_t i = layers - 1; i >= 0; i--)
 *    {
 *        const struct packet_layer *layer = packet_get_layer(pkt, i);
 *        printf("layer[%d]: type=%d, hdr_offset=%d, hdr_len=%d, pld_len=%d\n", i, layer->type, layer->hdr_offset, layer->hdr_len, layer->pld_len);
 *    }
 */

struct address
{
    uint8_t family; // AF_INET or AF_INET6
    union
    {
        struct in_addr v4;  /* network order */
        struct in6_addr v6; /* network order */
    } data;
};

static inline int packet_get_addr(const struct packet *pkt, struct address *src_addr, struct address *dst_addr)
{
    const struct ip *ip4_hdr = NULL;
    const struct ip6_hdr *ip6_hdr = NULL;
    const struct packet_layer *layer = NULL;
    int8_t num = packet_get_layers_number(pkt);
    for (int8_t i = num - 1; i >= 0; i--)
    {
        layer = packet_get_layer(pkt, i);
        if (layer->type & LAYER_TYPE_IPV4)
        {
            ip4_hdr = (const struct ip *)layer->hdr_ptr;
            if (src_addr != NULL)
            {
                src_addr->family = AF_INET;
                src_addr->data.v4.s_addr = ip4_hdr->ip_src.s_addr;
            }
            if (dst_addr != NULL)
            {
                dst_addr->family = AF_INET;
                dst_addr->data.v4.s_addr = ip4_hdr->ip_dst.s_addr;
            }
            return 0;
        }
        if (layer->type & LAYER_TYPE_IPV6)
        {
            ip6_hdr = (const struct ip6_hdr *)layer->hdr_ptr;
            if (src_addr != NULL)
            {
                src_addr->family = AF_INET6;
                src_addr->data.v6 = ip6_hdr->ip6_src;
            }
            if (dst_addr != NULL)
            {
                dst_addr->family = AF_INET6;
                dst_addr->data.v6 = ip6_hdr->ip6_dst;
            }
            return 0;
        }
    }

    return -1;
}

static inline int packet_get_port(const struct packet *pkt, uint16_t *src_port, uint16_t *dst_port)
{
    const struct tcphdr *tcp_hdr = NULL;
    const struct udphdr *udp_hdr = NULL;
    const struct packet_layer *layer = NULL;
    int8_t num = packet_get_layers_number(pkt);
    for (int8_t i = num - 1; i >= 0; i--)
    {
        layer = packet_get_layer(pkt, i);
        if (layer->type & LAYER_TYPE_TCP)
        {
            tcp_hdr = (const struct tcphdr *)layer->hdr_ptr;
            src_port != NULL ? *src_port = tcp_hdr->th_sport : 0;
            dst_port != NULL ? *dst_port = tcp_hdr->th_dport : 0;
            return 0;
        }
        if (layer->type & LAYER_TYPE_UDP)
        {
            udp_hdr = (const struct udphdr *)layer->hdr_ptr;
            src_port != NULL ? *src_port = udp_hdr->uh_sport : 0;
            dst_port != NULL ? *dst_port = udp_hdr->uh_dport : 0;
            return 0;
        }
    }

    return -1;
}

static inline int packet_get_ip_hdr(const struct packet *pkt, struct ip *hdr)
{
    const struct packet_layer *layer = NULL;
    int8_t num = packet_get_layers_number(pkt);
    for (int8_t i = num - 1; i >= 0; i--)
    {
        layer = packet_get_layer(pkt, i);
        if (layer->type & LAYER_TYPE_IPV4)
        {
            if (hdr != NULL)
            {
                memcpy(hdr, layer->hdr_ptr, sizeof(struct ip));
            }
            return 0;
        }
    }

    return -1;
}

static inline int packet_get_ip6_hdr(const struct packet *pkt, struct ip6_hdr *hdr)
{
    const struct packet_layer *layer = NULL;
    int8_t num = packet_get_layers_number(pkt);
    for (int8_t i = num - 1; i >= 0; i--)
    {
        layer = packet_get_layer(pkt, i);
        if (layer->type & LAYER_TYPE_IPV6)
        {
            if (hdr != NULL)
            {
                memcpy(hdr, layer->hdr_ptr, sizeof(struct ip6_hdr));
            }
            return 0;
        }
    }

    return -1;
}

static inline int packet_get_tcp_hdr(const struct packet *pkt, struct tcphdr *hdr)
{
    const struct packet_layer *layer = NULL;
    int8_t num = packet_get_layers_number(pkt);
    for (int8_t i = num - 1; i >= 0; i--)
    {
        layer = packet_get_layer(pkt, i);
        if (layer->type & LAYER_TYPE_TCP)
        {
            if (hdr != NULL)
            {
                memcpy(hdr, layer->hdr_ptr, sizeof(struct tcphdr));
            }
            return 0;
        }
    }

    return -1;
}

static inline int packet_get_udp_hdr(const struct packet *pkt, struct udphdr *hdr)
{
    const struct packet_layer *layer = NULL;
    int8_t num = packet_get_layers_number(pkt);
    for (int8_t i = num - 1; i >= 0; i--)
    {
        layer = packet_get_layer(pkt, i);
        if (layer->type & LAYER_TYPE_UDP)
        {
            if (hdr != NULL)
            {
                memcpy(hdr, layer->hdr_ptr, sizeof(struct udphdr));
            }
            return 0;
        }
    }

    return -1;
}

#ifdef __cplusplus
}
#endif

#endif