#pragma once

#include <gtest/gtest.h>

#ifdef __cplusplus
extern "C"
{
#endif

#include "packet_helper.h"
#include "ip_reassembly.h"
#include "packet_private.h"
#include "packet_dump.h"
#include "packet_parser.h"

static inline void packet_set_ipv4_src_addr(struct packet *pkt, uint32_t saddr)
{
    const struct layer_private *ipv4_layer = packet_get_innermost_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 layer_private *ipv6_layer = packet_get_innermost_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 layer_private *ipv6_layer = packet_get_innermost_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,
                              // IPv4 frag stat
                              uint64_t ip4_defrags_expected,
                              uint64_t ip4_defrags_succeed,
                              uint64_t ip4_defrags_failed_timeout,
                              uint64_t ip4_defrags_failed_invalid_length,
                              uint64_t ip4_defrags_failed_overlap,
                              uint64_t ip4_defrags_failed_too_many_frag,

                              uint64_t ip4_frags,
                              uint64_t ip4_frags_freed,
                              uint64_t ip4_frags_buffered,
                              uint64_t ip4_frags_bypass_no_buffer,
                              uint64_t ip4_frags_bypass_dup_fist_frag,
                              uint64_t ip4_frags_bypass_dup_last_frag,

                              // IPv6 frag stat
                              uint64_t ip6_defrags_expected,
                              uint64_t ip6_defrags_succeed,
                              uint64_t ip6_defrags_failed_timeout,
                              uint64_t ip6_defrags_failed_invalid_length,
                              uint64_t ip6_defrags_failed_overlap,
                              uint64_t ip6_defrags_failed_too_many_frag,

                              uint64_t ip6_frags,
                              uint64_t ip6_frags_freed,
                              uint64_t ip6_frags_buffered,
                              uint64_t ip6_frags_bypass_no_buffer,
                              uint64_t ip6_frags_bypass_dup_fist_frag,
                              uint64_t ip6_frags_bypass_dup_last_frag)
{
    EXPECT_TRUE(stat != NULL);

    EXPECT_TRUE(stat->ip4_defrags_expected == ip4_defrags_expected);
    EXPECT_TRUE(stat->ip4_defrags_succeed == ip4_defrags_succeed);
    EXPECT_TRUE(stat->ip4_defrags_failed_timeout == ip4_defrags_failed_timeout);
    EXPECT_TRUE(stat->ip4_defrags_failed_invalid_length == ip4_defrags_failed_invalid_length);
    EXPECT_TRUE(stat->ip4_defrags_failed_overlap == ip4_defrags_failed_overlap);
    EXPECT_TRUE(stat->ip4_defrags_failed_too_many_frag == ip4_defrags_failed_too_many_frag);

    EXPECT_TRUE(stat->ip4_frags == ip4_frags);
    EXPECT_TRUE(stat->ip4_frags_freed == ip4_frags_freed);
    EXPECT_TRUE(stat->ip4_frags_buffered == ip4_frags_buffered);
    EXPECT_TRUE(stat->ip4_frags_bypass_no_buffer == ip4_frags_bypass_no_buffer);
    EXPECT_TRUE(stat->ip4_frags_bypass_dup_fist_frag == ip4_frags_bypass_dup_fist_frag);
    EXPECT_TRUE(stat->ip4_frags_bypass_dup_last_frag == ip4_frags_bypass_dup_last_frag);

    EXPECT_TRUE(stat->ip6_defrags_expected == ip6_defrags_expected);
    EXPECT_TRUE(stat->ip6_defrags_succeed == ip6_defrags_succeed);
    EXPECT_TRUE(stat->ip6_defrags_failed_timeout == ip6_defrags_failed_timeout);
    EXPECT_TRUE(stat->ip6_defrags_failed_invalid_length == ip6_defrags_failed_invalid_length);
    EXPECT_TRUE(stat->ip6_defrags_failed_overlap == ip6_defrags_failed_overlap);
    EXPECT_TRUE(stat->ip6_defrags_failed_too_many_frag == ip6_defrags_failed_too_many_frag);

    EXPECT_TRUE(stat->ip6_frags == ip6_frags);
    EXPECT_TRUE(stat->ip6_frags_freed == ip6_frags_freed);
    EXPECT_TRUE(stat->ip6_frags_buffered == ip6_frags_buffered);
    EXPECT_TRUE(stat->ip6_frags_bypass_no_buffer == ip6_frags_bypass_no_buffer);
    EXPECT_TRUE(stat->ip6_frags_bypass_dup_fist_frag == ip6_frags_bypass_dup_fist_frag);
    EXPECT_TRUE(stat->ip6_frags_bypass_dup_last_frag == ip6_frags_bypass_dup_last_frag);
}

#ifdef __cplusplus
}
#endif
