#pragma once

#include <gtest/gtest.h>

#ifdef __cplusplus
extern "C"
{
#endif

#include "udp_utils.h"
#include "tcp_utils.h"
#include "ip4_utils.h"
#include "ip6_utils.h"
#include "ip_reassembly.h"
#include "packet_def.h"
#include "packet_dump.h"
#include "packet_utils.h"
#include "packet_parse.h"
#include "packet_layer.h"

static inline void packet_set_ipv4_src_addr(struct packet *pkt, uint32_t saddr)
{
    const struct raw_layer *ipv4_layer = packet_get_innermost_raw_layer(pkt, LAYER_PROTO_IPV4);
    EXPECT_TRUE(ipv4_layer);
    struct ip *hdr = (struct ip *)ipv4_layer->hdr_ptr;
    ip4_hdr_set_src_addr(hdr, saddr);
}

static inline void packet_set_ipv6_src_addr(struct packet *pkt, struct in6_addr saddr)
{
    const struct raw_layer *ipv6_layer = packet_get_innermost_raw_layer(pkt, LAYER_PROTO_IPV6);
    EXPECT_TRUE(ipv6_layer);
    struct ip6_hdr *hdr = (struct ip6_hdr *)ipv6_layer->hdr_ptr;
    ip6_hdr_set_src_in6_addr(hdr, saddr);
}

static inline void packet_set_ipv6_frag_offset(struct packet *pkt, uint16_t offset)
{
    const struct raw_layer *ipv6_layer = packet_get_innermost_raw_layer(pkt, LAYER_PROTO_IPV6);
    EXPECT_TRUE(ipv6_layer);
    struct ip6_hdr *hdr = (struct ip6_hdr *)ipv6_layer->hdr_ptr;
    struct ip6_frag *frag_hdr = ip6_hdr_get_frag_ext(hdr);
    EXPECT_TRUE(frag_hdr);
    ipv6_frag_set_offset(frag_hdr, offset);
}

static inline void check_stat(struct ip_reassembly_stat *stat,
                              uint64_t ip4_flow_find, uint64_t ip4_flow_add, uint64_t ip4_flow_del, uint64_t ip4_flow_timeout,
                              uint64_t ip4_flow_fail_no_space, uint64_t ip4_flow_fail_overlap, uint64_t ip4_flow_fail_many_frag,
                              uint64_t ip4_flow_fail_invalid_length, uint64_t ip4_flow_bypass_dup_fist_frag, uint64_t ip4_flow_bypass_dup_last_frag,
                              uint64_t ip6_flow_find, uint64_t ip6_flow_add, uint64_t ip6_flow_del, uint64_t ip6_flow_timeout,
                              uint64_t ip6_flow_fail_no_space, uint64_t ip6_flow_fail_overlap, uint64_t ip6_flow_fail_many_frag,
                              uint64_t ip6_flow_fail_invalid_length, uint64_t ip6_flow_bypass_dup_fist_frag, uint64_t ip6_flow_bypass_dup_last_frag)
{
    EXPECT_TRUE(stat != NULL);

    EXPECT_TRUE(stat->ip4_flow_find == ip4_flow_find);
    EXPECT_TRUE(stat->ip4_flow_add == ip4_flow_add);
    EXPECT_TRUE(stat->ip4_flow_del == ip4_flow_del);
    EXPECT_TRUE(stat->ip4_flow_timeout == ip4_flow_timeout);

    EXPECT_TRUE(stat->ip4_flow_fail_no_space == ip4_flow_fail_no_space);
    EXPECT_TRUE(stat->ip4_flow_fail_overlap == ip4_flow_fail_overlap);
    EXPECT_TRUE(stat->ip4_flow_fail_many_frag == ip4_flow_fail_many_frag);

    EXPECT_TRUE(stat->ip4_flow_fail_invalid_length == ip4_flow_fail_invalid_length);
    EXPECT_TRUE(stat->ip4_flow_bypass_dup_fist_frag == ip4_flow_bypass_dup_fist_frag);
    EXPECT_TRUE(stat->ip4_flow_bypass_dup_last_frag == ip4_flow_bypass_dup_last_frag);

    EXPECT_TRUE(stat->ip6_flow_find == ip6_flow_find);
    EXPECT_TRUE(stat->ip6_flow_add == ip6_flow_add);
    EXPECT_TRUE(stat->ip6_flow_del == ip6_flow_del);
    EXPECT_TRUE(stat->ip6_flow_timeout == ip6_flow_timeout);

    EXPECT_TRUE(stat->ip6_flow_fail_no_space == ip6_flow_fail_no_space);
    EXPECT_TRUE(stat->ip6_flow_fail_overlap == ip6_flow_fail_overlap);
    EXPECT_TRUE(stat->ip6_flow_fail_many_frag == ip6_flow_fail_many_frag);

    EXPECT_TRUE(stat->ip6_flow_fail_invalid_length == ip6_flow_fail_invalid_length);
    EXPECT_TRUE(stat->ip6_flow_bypass_dup_fist_frag == ip6_flow_bypass_dup_fist_frag);
    EXPECT_TRUE(stat->ip6_flow_bypass_dup_last_frag == ip6_flow_bypass_dup_last_frag);
}

#ifdef __cplusplus
}
#endif