#include <stdbool.h>
#include "udp_packet.h"
#include "../../utils.h"
#include "ipv4_packet.h"
#include "ethernet_frame.h"

#define ETH_UDP_HEADER_SIZE (8)

typedef struct {
    uint32_t sourceIpAddr;
    uint32_t destIpAddr;
    uint8_t zero;
    uint8_t protocol;
    uint16_t udpLength;
} UdpPseudoHeader;

static uint16_t udp_checksum(const UdpPseudoHeader *pseudoHeader, const uint8_t * hdr, uint32_t size) {
    uint32_t chksum = 0;

    // pseudoheader
    for (uint16_t i = 0; i < sizeof(UdpPseudoHeader) / 2; i++) {
        uint16_t field = ((uint16_t *) pseudoHeader)[i];
        field = htons(field);
        chksum += field;
    }

    // UDP header and hdr
    for (uint16_t i = 0; i < size / 2; i++) {
        if (i != 3) { // skip Checksum field
            uint16_t field = ((uint16_t *) hdr)[i];
            field = htons(field);
            chksum += field;
        }
    }

    // pad if packet size is odd
    if (size % 2) {
        chksum += ((uint16_t) (hdr[size - 1] << 8));
    }

    while (chksum >> 16) {
        chksum = (chksum & 0xFFFF) + (chksum >> 16);
    }

    uint16_t sum16 = ~(chksum & 0xFFFF);
    return sum16;
}

int parse_udp(const uint8_t *hdr, uint32_t size, PcktHeaderElement *pcktHdrLe, struct EthInterface_ *intf) {
    UdpProps *udpProps = HEADER_FETCH_PROPS(UdpProps, pcktHdrLe);
    FETCH_WORD_H2N_ADVANCE(&udpProps->SourcePort, hdr);
    FETCH_WORD_H2N_ADVANCE(&udpProps->DestinationPort, hdr);
    FETCH_WORD_H2N_ADVANCE(&udpProps->Length, hdr);
    FETCH_WORD_H2N_ADVANCE(&udpProps->Checksum, hdr);

    // common fields...
    udpProps->headerSize = ETH_UDP_HEADER_SIZE;
    udpProps->validityOK = (size == udpProps->Length); // TODO UDP checksum validation!
    udpProps->containedPacketClass = 0;

    return udpProps->validityOK ? 0 : -1;
}

void insert_udp_header(uint8_t *hdr, const PcktHeaderElement *headers) {
    uint8_t *hdrBegin = hdr;
    UdpProps *udpProps = (UdpProps *) &headers->props;
    FILL_WORD_H2N_ADVANCE(hdr, udpProps->SourcePort);
    FILL_WORD_H2N_ADVANCE(hdr, udpProps->DestinationPort);
    FILL_WORD_H2N_ADVANCE(hdr, udpProps->Length);

    // calculate checksum
    const IPv4Props *ipProps = (IPv4Props *) &headers->prev->props;
    UdpPseudoHeader ph = {ipProps->SourceIPAddr, ipProps->DestIPAddr, 0, ETH_UDP_PACKET_CLASS, htons(udpProps->Length)};
    udpProps->Checksum = udp_checksum(&ph, hdrBegin, udpProps->Length);
    FILL_WORD_H2N_ADVANCE(hdr, udpProps->Checksum);
}