#pragma once

#ifdef __cplusplus
extern "C"
{
#endif

#include <stdio.h>
#include <string.h>
#include <arpa/inet.h>
#define __FAVOR_BSD 1
#include <netinet/tcp.h>

/*
 * TCP 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
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |          Source Port          |       Destination Port        |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                        Sequence Number                        |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                    Acknowledgment Number                      |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |  Data |           |U|A|P|R|S|F|                               |
 *  | Offset| Reserved  |R|C|S|S|Y|I|            Window             |
 *  |       |           |G|K|H|T|N|N|                               |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |           Checksum            |         Urgent Pointer        |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                    Options                    |    Padding    |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 *  |                             data                              |
 *  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 */

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

static inline uint16_t tcp_hdr_get_src_port(const struct tcphdr *hdr)
{
    return ntohs(hdr->th_sport);
}

static inline uint16_t tcp_hdr_get_dst_port(const struct tcphdr *hdr)
{
    return ntohs(hdr->th_dport);
}

static inline uint32_t tcp_hdr_get_seq(const struct tcphdr *hdr)
{
    return ntohl(hdr->th_seq);
}

static inline uint32_t tcp_hdr_get_ack(const struct tcphdr *hdr)
{
    return ntohl(hdr->th_ack);
}

static inline uint8_t tcp_hdr_get_hdr_len(const struct tcphdr *hdr)
{
    return hdr->th_off << 2;
}

static inline uint8_t tcp_hdr_get_flags(const struct tcphdr *hdr)
{
    return hdr->th_flags;
}

static inline bool tcp_hdr_get_urg_flag(const struct tcphdr *hdr)
{
    return hdr->th_flags & TH_URG;
}

static inline bool tcp_hdr_get_ack_flag(const struct tcphdr *hdr)
{
    return hdr->th_flags & TH_ACK;
}

static inline bool tcp_hdr_get_push_flag(const struct tcphdr *hdr)
{
    return hdr->th_flags & TH_PUSH;
}

static inline bool tcp_hdr_get_rst_flag(const struct tcphdr *hdr)
{
    return hdr->th_flags & TH_RST;
}

static inline bool tcp_hdr_get_syn_flag(const struct tcphdr *hdr)
{
    return hdr->th_flags & TH_SYN;
}

static inline bool tcp_hdr_get_fin_flag(const struct tcphdr *hdr)
{
    return hdr->th_flags & TH_FIN;
}

static inline uint16_t tcp_hdr_get_window(const struct tcphdr *hdr)
{
    return ntohs(hdr->th_win);
}

static inline uint16_t tcp_hdr_get_checksum(const struct tcphdr *hdr)
{
    return ntohs(hdr->th_sum);
}

static inline uint16_t tcp_hdr_get_urg_ptr(const struct tcphdr *hdr)
{
    return ntohs(hdr->th_urp);
}

static inline uint16_t tcp_hdr_get_opt_len(const struct tcphdr *hdr)
{
    return tcp_hdr_get_hdr_len(hdr) - sizeof(struct tcphdr);
}

static inline const char *tcp_hdr_get_opt_data(const struct tcphdr *hdr)
{
    if (tcp_hdr_get_opt_len(hdr) == 0)
    {
        return NULL;
    }

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

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

static inline void tcp_hdr_set_src_port(struct tcphdr *hdr, uint16_t port)
{
    hdr->th_sport = htons(port);
}

static inline void tcp_hdr_set_dst_port(struct tcphdr *hdr, uint16_t port)
{
    hdr->th_dport = htons(port);
}

static inline void tcp_hdr_set_seq(struct tcphdr *hdr, uint32_t seq)
{
    hdr->th_seq = htonl(seq);
}

static inline void tcp_hdr_set_ack(struct tcphdr *hdr, uint32_t ack)
{
    hdr->th_ack = htonl(ack);
}

static inline void tcp_hdr_set_hdr_len(struct tcphdr *hdr, uint8_t offset)
{
    hdr->th_off = offset >> 2;
}

static inline void tcp_hdr_set_flags(struct tcphdr *hdr, uint8_t flags)
{
    hdr->th_flags = flags;
}

static inline void tcp_hdr_set_urg_flag(struct tcphdr *hdr, bool flag)
{
    if (flag)
    {
        hdr->th_flags |= TH_URG;
    }
    else
    {
        hdr->th_flags &= ~TH_URG;
    }
}

static inline void tcp_hdr_set_ack_flag(struct tcphdr *hdr, bool flag)
{
    if (flag)
    {
        hdr->th_flags |= TH_ACK;
    }
    else
    {
        hdr->th_flags &= ~TH_ACK;
    }
}

static inline void tcp_hdr_set_push_flag(struct tcphdr *hdr, bool flag)
{
    if (flag)
    {
        hdr->th_flags |= TH_PUSH;
    }
    else
    {
        hdr->th_flags &= ~TH_PUSH;
    }
}

static inline void tcp_hdr_set_rst_flag(struct tcphdr *hdr, bool flag)
{
    if (flag)
    {
        hdr->th_flags |= TH_RST;
    }
    else
    {
        hdr->th_flags &= ~TH_RST;
    }
}

static inline void tcp_hdr_set_syn_flag(struct tcphdr *hdr, bool flag)
{
    if (flag)
    {
        hdr->th_flags |= TH_SYN;
    }
    else
    {
        hdr->th_flags &= ~TH_SYN;
    }
}

static inline void tcp_hdr_set_fin_flag(struct tcphdr *hdr, bool flag)
{
    if (flag)
    {
        hdr->th_flags |= TH_FIN;
    }
    else
    {
        hdr->th_flags &= ~TH_FIN;
    }
}

static inline void tcp_hdr_set_window(struct tcphdr *hdr, uint16_t window)
{
    hdr->th_win = htons(window);
}

static inline void tcp_hdr_set_checksum(struct tcphdr *hdr, uint16_t checksum)
{
    hdr->th_sum = htons(checksum);
}

static inline void tcp_hdr_set_urg_ptr(struct tcphdr *hdr, uint16_t ptr)
{
    hdr->th_urp = htons(ptr);
}

static inline void tcp_hdr_set_opt_len(struct tcphdr *hdr, uint16_t len)
{
    hdr->th_off = (sizeof(struct tcphdr) + len) >> 2;
}

// must be called after tcp_hdr_set_opt_len
static inline void tcp_hdr_set_opt_data(struct tcphdr *hdr, const char *ptr)
{
    memcpy((char *)hdr + sizeof(struct tcphdr), ptr, tcp_hdr_get_opt_len(hdr));
}

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

static inline int tcp_hdr_to_str(const struct tcphdr *hdr, char *buf, size_t size)
{
    memset(buf, 0, size);
    return snprintf(buf, size, "TCP: src_port=%u dst_port=%u seq=%u ack=%u hdr_len=%u flags=0x%02x(%s%s%s%s%s%s) window=%u checksum=0x%x urg_ptr=%u opt_len=%u",
                    tcp_hdr_get_src_port(hdr), tcp_hdr_get_dst_port(hdr),
                    tcp_hdr_get_seq(hdr), tcp_hdr_get_ack(hdr),
                    tcp_hdr_get_hdr_len(hdr), tcp_hdr_get_flags(hdr),
                    tcp_hdr_get_urg_flag(hdr) ? "URG " : "",
                    tcp_hdr_get_ack_flag(hdr) ? "ACK " : "",
                    tcp_hdr_get_push_flag(hdr) ? "PSH " : "",
                    tcp_hdr_get_rst_flag(hdr) ? "RST " : "",
                    tcp_hdr_get_syn_flag(hdr) ? "SYN " : "",
                    tcp_hdr_get_fin_flag(hdr) ? "FIN " : "",
                    tcp_hdr_get_window(hdr), tcp_hdr_get_checksum(hdr),
                    tcp_hdr_get_urg_ptr(hdr), tcp_hdr_get_opt_len(hdr));
}

#ifdef __cplusplus
}
#endif