#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_proto
{
    LAYER_PROTO_NONE = 0,

    // L2 -- data link layer
    LAYER_PROTO_ETHER = 1,
    LAYER_PROTO_PWETH = 2,
    LAYER_PROTO_PPP = 3,
    LAYER_PROTO_L2TP = 4,

    // L2 -- tunnel
    LAYER_PROTO_VLAN = 21,
    LAYER_PROTO_PPPOE = 22,
    LAYER_PROTO_MPLS = 23,

    // L3 -- network layer
    LAYER_PROTO_IPV4 = 31,
    LAYER_PROTO_IPV6 = 32,
    LAYER_PROTO_IPAH = 33,

    // L3 -- tunnel
    LAYER_PROTO_GRE = 41,

    // L4 -- transport layer
    LAYER_PROTO_UDP = 51,
    LAYER_PROTO_TCP = 52,
    LAYER_PROTO_ICMP = 53,
    LAYER_PROTO_ICMP6 = 54,

    // L4 -- tunnel
    LAYER_PROTO_VXLAN = 61,
    LAYER_PROTO_GTPV1_U = 62,
};

struct raw_layer
{
    enum layer_proto 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);

int packet_get_layer_count(const struct packet *pkt);
const struct raw_layer *packet_get_raw_layer(const struct packet *pkt, int 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);

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 raw_layer *layer = NULL;
    int num = packet_get_layer_count(pkt);
    for (int i = num - 1; i >= 0; i--)
    {
        layer = packet_get_raw_layer(pkt, i);
        if (layer->type == LAYER_PROTO_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_PROTO_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 raw_layer *layer = NULL;
    int num = packet_get_layer_count(pkt);
    for (int i = num - 1; i >= 0; i--)
    {
        layer = packet_get_raw_layer(pkt, i);
        if (layer->type == LAYER_PROTO_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_PROTO_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;
}

#ifdef __cplusplus
}
#endif

#endif