#pragma once

#ifdef __cplusplus
extern "C"
{
#endif

#include <stdio.h>
#include <string.h>
#include <arpa/inet.h>
#include <netinet/ip.h>

/*
 * Internet Header Format
 *
 *   0                   1                   2                   3
 *   0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |Version|  IHL  |Type of Service|          Total Length         |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |         Identification        |Flags|      Fragment Offset    |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |  Time to Live |    Protocol   |         Header Checksum       |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                       Source Address                          |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                    Destination Address                        |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                    Options                    |    Padding    |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 */

/******************************************************************************
 * get
 ******************************************************************************/

static inline uint8_t ip4_hdr_get_version(const struct ip *hdr)
{
    return hdr->ip_v;
}

// IP Options are included in the hdr_len field
static inline uint8_t ip4_hdr_get_hdr_len(const struct ip *hdr)
{
    return hdr->ip_hl << 2;
}

static inline uint8_t ip4_hdr_get_tos(const struct ip *hdr)
{
    return hdr->ip_tos;
}

static inline uint16_t ip4_hdr_get_total_len(const struct ip *hdr)
{
    return ntohs(hdr->ip_len);
}

static inline uint16_t ip4_hdr_get_ipid(const struct ip *hdr)
{
    return ntohs(hdr->ip_id);
}

static inline uint8_t ip4_hdr_get_flags(const struct ip *hdr)
{
    return (ntohs(hdr->ip_off) & (~IP_OFFMASK)) >> 13;
}

static inline bool ip4_hdr_get_rf_flag(const struct ip *hdr)
{
    return (ntohs(hdr->ip_off) & IP_RF) != 0;
}

static inline bool ip4_hdr_get_df_flag(const struct ip *hdr)
{
    return (ntohs(hdr->ip_off) & IP_DF) != 0;
}

static inline bool ip4_hdr_get_mf_flag(const struct ip *hdr)
{
    return (ntohs(hdr->ip_off) & IP_MF) != 0;
}

static inline uint16_t ip4_hdr_get_frag_offset(const struct ip *hdr)
{
    return (ntohs(hdr->ip_off) & IP_OFFMASK) << 3;
}

static inline uint8_t ip4_hdr_get_ttl(const struct ip *hdr)
{
    return hdr->ip_ttl;
}

static inline uint8_t ip4_hdr_get_proto(const struct ip *hdr)
{
    return hdr->ip_p;
}

static inline uint16_t ip4_hdr_get_checksum(const struct ip *hdr)
{
    return ntohs(hdr->ip_sum);
}

static inline uint32_t ip4_hdr_get_src_addr(const struct ip *hdr)
{
    return ntohl(hdr->ip_src.s_addr);
}

static inline uint32_t ip4_hdr_get_dst_addr(const struct ip *hdr)
{
    return ntohl(hdr->ip_dst.s_addr);
}

static inline struct in_addr ip4_hdr_get_src_in_addr(const struct ip *hdr)
{
    return hdr->ip_src;
}

static inline struct in_addr ip4_hdr_get_dst_in_addr(const struct ip *hdr)
{
    return hdr->ip_dst;
}

static inline uint8_t ip4_hdr_get_opt_len(const struct ip *hdr)
{
    return ip4_hdr_get_hdr_len(hdr) - sizeof(struct ip);
}

static inline const char *ip4_hdr_get_opt_data(const struct ip *hdr)
{
    if (ip4_hdr_get_opt_len(hdr) == 0)
    {
        return NULL;
    }

    return (const char *)hdr + sizeof(struct ip);
}

/******************************************************************************
 * set
 ******************************************************************************/

static inline void ip4_hdr_set_version(struct ip *hdr, uint8_t version)
{
    hdr->ip_v = version;
}

static inline void ip4_hdr_set_hdr_len(struct ip *hdr, uint8_t hdr_len)
{
    hdr->ip_hl = hdr_len >> 2;
}

static inline void ip4_hdr_set_tos(struct ip *hdr, uint8_t tos)
{
    hdr->ip_tos = tos;
}

static inline void ip4_hdr_set_total_len(struct ip *hdr, uint16_t total_len)
{
    hdr->ip_len = htons(total_len);
}

static inline void ip4_hdr_set_ipid(struct ip *hdr, uint16_t ipid)
{
    hdr->ip_id = htons(ipid);
}

static inline void ip4_hdr_set_flags(struct ip *hdr, uint8_t flags)
{
    hdr->ip_off = htons((flags << 13) | (ntohs(hdr->ip_off) & IP_OFFMASK));
}

static inline void ip4_hdr_set_rf_flag(struct ip *hdr, bool flag)
{
    if (flag)
    {
        hdr->ip_off = htons(ntohs(hdr->ip_off) | IP_RF);
    }
    else
    {
        hdr->ip_off = htons(ntohs(hdr->ip_off) & ~IP_RF);
    }
}

static inline void ip4_hdr_set_df_flag(struct ip *hdr, bool flag)
{
    if (flag)
    {
        hdr->ip_off = htons(ntohs(hdr->ip_off) | IP_DF);
    }
    else
    {
        hdr->ip_off = htons(ntohs(hdr->ip_off) & ~IP_DF);
    }
}

static inline void ip4_hdr_set_mf_flag(struct ip *hdr, bool flag)
{
    if (flag)
    {
        hdr->ip_off = htons(ntohs(hdr->ip_off) | IP_MF);
    }
    else
    {
        hdr->ip_off = htons(ntohs(hdr->ip_off) & ~IP_MF);
    }
}

static inline void ip4_hdr_set_frag_offset(struct ip *hdr, uint16_t frag_offset)
{
    hdr->ip_off = htons((frag_offset >> 3) | (ntohs(hdr->ip_off) & ~IP_OFFMASK));
}

static inline void ip4_hdr_set_ttl(struct ip *hdr, uint8_t ttl)
{
    hdr->ip_ttl = ttl;
}

static inline void ip4_hdr_set_protocol(struct ip *hdr, uint8_t protocol)
{
    hdr->ip_p = protocol;
}

static inline void ip4_hdr_set_checksum(struct ip *hdr, uint16_t checksum)
{
    hdr->ip_sum = htons(checksum);
}

static inline void ip4_hdr_set_src_addr(struct ip *hdr, uint32_t saddr)
{
    hdr->ip_src.s_addr = htonl(saddr);
}

static inline void ip4_hdr_set_dst_addr(struct ip *hdr, uint32_t daddr)
{
    hdr->ip_dst.s_addr = htonl(daddr);
}

static inline void ip4_hdr_set_src_in_addr(struct ip *hdr, struct in_addr saddr)
{
    hdr->ip_src = saddr;
}

static inline void ip4_hdr_set_dst_in_addr(struct ip *hdr, struct in_addr daddr)
{
    hdr->ip_dst = daddr;
}

static inline void ip4_hdr_set_opt_len(struct ip *hdr, uint8_t opt_len)
{
    ip4_hdr_set_hdr_len(hdr, opt_len + sizeof(struct ip));
}

// must be called after ip4_hdr_set_opt_len
static inline void ip4_hdr_set_opt_data(struct ip *hdr, const char *opt_data)
{
    memcpy((char *)hdr + sizeof(struct ip), opt_data, ip4_hdr_get_opt_len(hdr));
}

/******************************************************************************
 * identify
 ******************************************************************************/

// /usr/include/netinet/in.h
static inline int is_ip_proto(uint16_t proto)
{
    switch (proto)
    {
    case IPPROTO_IP:
    case IPPROTO_ICMP:
    case IPPROTO_IGMP:
    case IPPROTO_IPIP:
    case IPPROTO_TCP:
    case IPPROTO_EGP:
    case IPPROTO_PUP:
    case IPPROTO_UDP:
    case IPPROTO_IDP:
    case IPPROTO_TP:
    case IPPROTO_DCCP:
    case IPPROTO_IPV6:
    case IPPROTO_RSVP:
    case IPPROTO_GRE:
    case IPPROTO_ESP:
    case IPPROTO_AH:
    case IPPROTO_MTP:
    case IPPROTO_BEETPH:
    case IPPROTO_ENCAP:
    case IPPROTO_PIM:
    case IPPROTO_COMP:
    case IPPROTO_SCTP:
    case IPPROTO_UDPLITE:
    case IPPROTO_MPLS:
    case IPPROTO_ETHERNET:
    case IPPROTO_RAW:
    case IPPROTO_MPTCP:
        return 1;
    default:
        return 0;
    }
}

/******************************************************************************
 * print
 ******************************************************************************/

static inline const char *ip_proto_to_str(uint16_t proto)
{
    switch (proto)
    {
    case IPPROTO_IP:
        return "IPPROTO_IP";
    case IPPROTO_ICMP:
        return "IPPROTO_ICMP";
    case IPPROTO_IGMP:
        return "IPPROTO_IGMP";
    case IPPROTO_IPIP:
        return "IPPROTO_IPIP";
    case IPPROTO_TCP:
        return "IPPROTO_TCP";
    case IPPROTO_EGP:
        return "IPPROTO_EGP";
    case IPPROTO_PUP:
        return "IPPROTO_PUP";
    case IPPROTO_UDP:
        return "IPPROTO_UDP";
    case IPPROTO_IDP:
        return "IPPROTO_IDP";
    case IPPROTO_TP:
        return "IPPROTO_TP";
    case IPPROTO_DCCP:
        return "IPPROTO_DCCP";
    case IPPROTO_IPV6:
        return "IPPROTO_IPV6";
    case IPPROTO_RSVP:
        return "IPPROTO_RSVP";
    case IPPROTO_GRE:
        return "IPPROTO_GRE";
    case IPPROTO_ESP:
        return "IPPROTO_ESP";
    case IPPROTO_AH:
        return "IPPROTO_AH";
    case IPPROTO_MTP:
        return "IPPROTO_MTP";
    case IPPROTO_BEETPH:
        return "IPPROTO_BEETPH";
    case IPPROTO_ENCAP:
        return "IPPROTO_ENCAP";
    case IPPROTO_PIM:
        return "IPPROTO_PIM";
    case IPPROTO_COMP:
        return "IPPROTO_COMP";
    case IPPROTO_SCTP:
        return "IPPROTO_SCTP";
    case IPPROTO_UDPLITE:
        return "IPPROTO_UDPLITE";
    case IPPROTO_MPLS:
        return "IPPROTO_MPLS";
    case IPPROTO_ETHERNET:
        return "IPPROTO_ETHERNET";
    case IPPROTO_RAW:
        return "IPPROTO_RAW";
    case IPPROTO_MPTCP:
        return "IPPROTO_MPTCP";
    default:
        return "IPPROTO_UNKNOWN";
    }
}

static inline int ip4_hdr_to_str(const struct ip *hdr, char *buf, size_t size)
{
    memset(buf, 0, size);
    char src_addr_str[INET6_ADDRSTRLEN] = {0};
    char dst_addr_str[INET6_ADDRSTRLEN] = {0};

    uint16_t proto = ip4_hdr_get_proto(hdr);
    struct in_addr src_addr = ip4_hdr_get_src_in_addr(hdr);
    struct in_addr dst_addr = ip4_hdr_get_dst_in_addr(hdr);
    inet_ntop(AF_INET, &src_addr, src_addr_str, sizeof(src_addr_str));
    inet_ntop(AF_INET, &dst_addr, dst_addr_str, sizeof(dst_addr_str));

    return snprintf(buf, size, "IPv4: version=%u hdr_len=%u tos=%u total_len=%u ipid=%u flags=%u(rf=%u df=%u mf=%u) frag_offset=%u ttl=%u proto=%s checksum=0x%x src_addr=%s dst_addr=%s opt_len=%u",
                    ip4_hdr_get_version(hdr), ip4_hdr_get_hdr_len(hdr), ip4_hdr_get_tos(hdr),
                    ip4_hdr_get_total_len(hdr), ip4_hdr_get_ipid(hdr), ip4_hdr_get_flags(hdr),
                    ip4_hdr_get_rf_flag(hdr), ip4_hdr_get_df_flag(hdr), ip4_hdr_get_mf_flag(hdr),
                    ip4_hdr_get_frag_offset(hdr), ip4_hdr_get_ttl(hdr), ip_proto_to_str(proto),
                    ip4_hdr_get_checksum(hdr), src_addr_str, dst_addr_str, ip4_hdr_get_opt_len(hdr));
}

#ifdef __cplusplus
}
#endif