#include <string.h>
#include <netinet/ip.h>
#include <netinet/in.h>
#include <netinet/ether.h>

#include "vxlan.h"

#define CHECKSUM_CARRY(x) (x = (x >> 16) + (x & 0xffff), (~(x + (x >> 16)) & 0xffff))
static inline int checksum(uint16_t *data, int len)
{
    int sum = 0;
    int nleft = len;
    uint16_t ans = 0;
    uint16_t *w = data;

    while (nleft > 1)
    {
        sum += *w++;
        nleft -= 2;
    }

    if (nleft == 1)
    {
        *(char *)(&ans) = *(char *)w;
        sum += ans;
    }

    return sum;
}

static inline void eth_header_encode(struct ethhdr *eth_hdr, const u_char src_mac[], const u_char dst_mac[], uint16_t proto)
{
    memcpy(eth_hdr->h_source, src_mac, ETH_ALEN);
    memcpy(eth_hdr->h_dest, dst_mac, ETH_ALEN);
    eth_hdr->h_proto = htons(proto);
}

static inline void ip_header_encode(struct ip *ip_hdr, const in_addr_t src_ip, const in_addr_t dst_ip, uint16_t ipid, uint8_t proto, uint16_t pld_len)
{
    ip_hdr->ip_v = 4;                                    /* version 4 */
    ip_hdr->ip_hl = 5;                                   /* 20 byte header */
    ip_hdr->ip_tos = 0;                                  /* IP tos */
    ip_hdr->ip_len = htons(sizeof(struct ip) + pld_len); /* total length */
    ip_hdr->ip_id = htons(ipid);                         /* IP ID */
    ip_hdr->ip_off = htons(0);                           /* fragmentation flags */
    ip_hdr->ip_ttl = 80;                                 /* time to live */
    ip_hdr->ip_p = proto;                                /* transport protocol */
    ip_hdr->ip_sum = 0;                                  /* do this later */
    ip_hdr->ip_src.s_addr = src_ip;
    ip_hdr->ip_dst.s_addr = dst_ip;

    int sum = checksum((uint16_t *)ip_hdr, 20);
    ip_hdr->ip_sum = CHECKSUM_CARRY(sum);
}

static inline void udp_header_encode(struct udphdr *udp_hdr, uint16_t udp_sport, uint16_t udp_dport, uint16_t pld_len)
{
    int udp_hlen = sizeof(struct udphdr) + pld_len;

    udp_hdr->uh_sport = htons(udp_sport);
    udp_hdr->uh_dport = htons(udp_dport);
    udp_hdr->uh_ulen = htons(udp_hlen);
    /*
     * UDP Checksum: It SHOULD be transmitted as zero
     * https://datatracker.ietf.org/doc/html/rfc7348#section-5
     */
    udp_hdr->uh_sum = 0;
}

// return  0 : success
// return -1 : error
int vxlan_frame_decode(struct vxlan_hdr **vxlan_hdr, const char *data, uint16_t len)
{
    if (len < VXLAN_FRAME_HDR_LEN)
    {
        return -1;
    }

    struct ethhdr *eth_hdr = (struct ethhdr *)data;
    if (eth_hdr->h_proto != htons(ETH_P_IP))
    {
        return -1;
    }

    struct ip *ip_hdr = (struct ip *)((char *)eth_hdr + sizeof(struct ethhdr));
    if (ip_hdr->ip_p != IPPROTO_UDP)
    {
        return -1;
    }

    struct udphdr *udp_hdr = (struct udphdr *)((char *)ip_hdr + sizeof(struct ip));
    if (udp_hdr->uh_dport != htons(4789))
    {
        return -1;
    }

    *vxlan_hdr = (struct vxlan_hdr *)((char *)udp_hdr + sizeof(struct udphdr));

    return 0;
}

void vxlan_frame_encode(char *buff,
                        const u_char eth_src_mac[], const u_char eth_dst_mac[],
                        const in_addr_t ip_src_addr, const in_addr_t ip_dst_addr, uint16_t ip_id,
                        uint16_t udp_src_port, uint16_t udp_pld_len,
                        uint8_t vni_opt_dir, uint8_t vni_opt_traffic, uint8_t vni_opt_sf_index)
{
    struct ethhdr *eth_hdr = (struct ethhdr *)buff;
    struct ip *ip_hdr = (struct ip *)((char *)eth_hdr + sizeof(struct ethhdr));
    struct udphdr *udp_hdr = (struct udphdr *)((char *)ip_hdr + sizeof(struct ip));
    struct vxlan_hdr *vxlan_hdr = (struct vxlan_hdr *)((char *)udp_hdr + sizeof(struct udphdr));

    // MUST be set to zero
    memset(vxlan_hdr, 0, sizeof(struct vxlan_hdr));
    vxlan_hdr->flags = VXLAN_FLAGS;
    vxlan_set_opt(vxlan_hdr, VNI_OPT_DIR, vni_opt_dir);
    vxlan_set_opt(vxlan_hdr, VNI_OPT_TRAFFIC, vni_opt_traffic);
    vxlan_set_opt(vxlan_hdr, VNI_OPT_SFINDEX, vni_opt_sf_index);

    eth_header_encode(eth_hdr, eth_src_mac, eth_dst_mac, ETH_P_IP);
    ip_header_encode(ip_hdr, ip_src_addr, ip_dst_addr, ip_id, IPPROTO_UDP, sizeof(struct udphdr) + sizeof(struct vxlan_hdr) + udp_pld_len);
    udp_header_encode(udp_hdr, udp_src_port, VXLAN_DST_PORT, sizeof(struct vxlan_hdr) + udp_pld_len);
}

uint16_t calculate_vxlan_source_port(struct four_tuple *innermost_tuple4)
{
    /*
     * When calculating the UDP source port number in this manner, it
     * is RECOMMENDED that the value be in the dynamic/private port
     * range 49152-65535 [RFC6335].
     */
    uint64_t hash = four_tuple_hash(innermost_tuple4);
    uint16_t port = (uint16_t)(hash % (65535 - 49152) + 49152);

    return port;
}