#include "dablooms.h"
#include "packet_dabloom.h"
#include "packet_internal.h"

struct packet_key
{
    // TCP or UDP
    uint32_t seq;         // network order
    uint32_t ack;         // network order
    uint16_t src_port;    // network order
    uint16_t dst_port;    // network order
    uint16_t l4_checksum; // network order

    // IPv4
    uint16_t ip_id;    // network order
    uint32_t src_addr; // network order
    uint32_t dst_addr; // network order
} __attribute__((__packed__));

struct packet_dabloom
{
    struct expiry_dablooms_handle *dablooms;
};

/******************************************************************************
 * Private API
 ******************************************************************************/

// return  0: success
// reutrn -1: error
static inline int packet_key_get(const struct packet *packet, struct packet_key *key)
{
    const struct layer_internal *ip_layer = packet_get_innermost_layer(packet, LAYER_PROTO_IPV4);
    const struct layer_internal *tcp_layer = packet_get_innermost_layer(packet, LAYER_PROTO_TCP);
    const struct layer_internal *udp_layer = packet_get_innermost_layer(packet, LAYER_PROTO_UDP);
    if (ip_layer == NULL || (tcp_layer == NULL && udp_layer == NULL))
    {
        return -1;
    }

    memset(key, 0, sizeof(struct packet_key));

    const struct ip *iphdr = (const struct ip *)ip_layer->hdr_ptr;
    key->ip_id = iphdr->ip_id;
    key->src_addr = iphdr->ip_src.s_addr;
    key->dst_addr = iphdr->ip_dst.s_addr;

    if (tcp_layer)
    {
        const struct tcphdr *tcphdr = (const struct tcphdr *)tcp_layer->hdr_ptr;
        key->seq = tcphdr->th_seq;
        key->ack = tcphdr->th_ack;
        key->src_port = tcphdr->th_sport;
        key->dst_port = tcphdr->th_dport;
        key->l4_checksum = tcphdr->th_sum;
    }
    else
    {
        const struct udphdr *udphdr = (const struct udphdr *)udp_layer->hdr_ptr;
        key->src_port = udphdr->uh_sport;
        key->dst_port = udphdr->uh_dport;
        key->l4_checksum = udphdr->uh_sum;
    }

    return 0;
}

/******************************************************************************
 * Public API
 ******************************************************************************/

struct packet_dabloom *packet_dabloom_new(uint32_t capacity, uint32_t timeout, double error_rate, uint64_t now)
{
    struct packet_dabloom *pkt_dab = (struct packet_dabloom *)calloc(1, sizeof(struct packet_dabloom));
    if (pkt_dab == NULL)
    {
        return NULL;
    }

    pkt_dab->dablooms = expiry_dablooms_new(capacity, error_rate, now, timeout);
    if (pkt_dab->dablooms == NULL)
    {
        free(pkt_dab);
        return NULL;
    }

    return pkt_dab;
}

void packet_dabloom_free(struct packet_dabloom *pkt_dab)
{
    if (pkt_dab)
    {
        if (pkt_dab->dablooms)
        {
            expiry_dablooms_free(pkt_dab->dablooms);
            pkt_dab->dablooms = NULL;
        }
        free(pkt_dab);
        pkt_dab = NULL;
    }
}

// return 1: found
// reutrn 0: no found
int packet_dabloom_lookup(struct packet_dabloom *pkt_dab, const struct packet *packet, uint64_t now)
{
    struct packet_key key;
    if (packet_key_get(packet, &key) == -1)
    {
        return 0;
    }

    if (expiry_dablooms_search(pkt_dab->dablooms, (const char *)&key, sizeof(struct packet_key), now) == 1)
    {
        return 1;
    }

    return 0;
}

void packet_dabloom_add(struct packet_dabloom *pkt_dab, const struct packet *packet, uint64_t now)
{
    struct packet_key key;
    if (packet_key_get(packet, &key) == -1)
    {
        return;
    }

    expiry_dablooms_add(pkt_dab->dablooms, (const char *)&key, sizeof(struct packet_key), now);
}