//
// Created by epagris on 2022.11.03..
//

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

#include "../../eth_interface.h"

#define ETH_IP_HEADER_LENGTH (20)

static bool check_ipv4_validity(const uint8_t *hdr, const IPv4Props *ipProps, bool dontVerifyChecksum) {
    bool valid =
            (ipProps->Version == 4) &&
            (ipProps->IHL == (ETH_IP_HEADER_LENGTH / 4)) &&
            (dontVerifyChecksum || (chksum(hdr, ETH_IPv4_HEADER_SIZE, false) == 0)); /*&&
            (ipProps->FragmentOffset == 0) && !(ipProps->Flags & 0x02 << 4); // TODO: discard if fragmented*/
    return valid;
}

int parse_ipv4(const uint8_t *hdr, uint32_t size, PcktHeaderElement *pcktHdrLe, struct EthInterface_ *intf, PcktProcFnPassbackData * pb) {
    const uint8_t *hdrBegin = hdr;
    IPv4Props *ipProps = HEADER_FETCH_PROPS(IPv4Props, pcktHdrLe);
    uint8_t version_length;
    FETCH_ADVANCE(&version_length, hdr, 1);
    ipProps->Version = (version_length) >> 4;
    ipProps->IHL = (version_length) & 0x0F;
    FETCH_ADVANCE(&ipProps->DSF, hdr, 1);
    FETCH_WORD_H2N_ADVANCE(&ipProps->TotalLength, hdr);
    FETCH_WORD_H2N_ADVANCE(&ipProps->Identification, hdr);
    uint16_t flags_fragOffset;
    FETCH_WORD_H2N_ADVANCE(&flags_fragOffset, hdr);
    ipProps->Flags = (flags_fragOffset >> 13) & 0x07;
    ipProps->FragmentOffset = (flags_fragOffset & ~(0x07 << 13)) << 3;
    FETCH_ADVANCE(&ipProps->TTL, hdr, 1);
    FETCH_ADVANCE(&ipProps->Protocol, hdr, 1);
    FETCH_WORD_H2N_ADVANCE(&ipProps->HeaderChecksum, hdr);
    FETCH_ADVANCE(&ipProps->SourceIPAddr, hdr, 4); // NO network -> host byte order conversion!
    FETCH_ADVANCE(&ipProps->DestIPAddr, hdr, 4);

    // fill-in common packet header fields
    ipProps->validityOK = check_ipv4_validity(hdrBegin, ipProps, (intf->capabilities & ETHINF_CAP_RX_IPCHKSUM_OFFLOAD)); // don't verify checksum if done in hardware
    ipProps->containedPacketClass = ipProps->Protocol;
    ipProps->headerSize = ETH_IP_HEADER_LENGTH;

    // if packet is valid, learn MAC-IP mapping
    if (ipProps->validityOK) {
        EthernetProps *ethProps = HEADER_FETCH_PROPS(EthernetProps, pcktHdrLe->prev); // fetch Ethernet header
        ArpEntry entry; // fill entry
        memcpy(&entry.eth, ethProps->sourceAddr, ETH_HW_ADDR_LEN);
        entry.ip = ipProps->SourceIPAddr;
        arpc_learn(intf->arpc, &entry); // learn new assignment
    }

    // check if packet has padding at the end, then strip it
    int ret = 0;
    if (size > ipProps->TotalLength) {
        pb->i = (int32_t)size - (int32_t)ipProps->TotalLength; // store the stripping length
        ret |= PROC_FN_FLAG_STRIP_PAD;
    }

    // check whether it is a fragmented packet
    bool isAFragment = (ipProps->FragmentOffset != 0) || (ipProps->Flags & IP_FLAG_MF);
    if (isAFragment) {
        const uint8_t * payloadStart = hdr;
        ipra_input(intf->ipra, ipProps, payloadStart, size - ETH_IP_HEADER_LENGTH); // input fragment into assembler
        uint8_t * raPload;
        uint16_t raSize;
        bool raOK = ipra_try_reassemble(intf->ipra, ipProps->Identification, &raPload, &raSize, pcktHdrLe);
        if (raOK) {
            pb->p = raPload;
            pb->u = raSize;
            pb->b = true; // MRD!
            ret |= PROC_FN_RET_REPRST;
            return ret; // replace buffer and restart processing (with the assembled packet)
        } else {
            ipProps->containedPacketClass = 0; // no contained packet if not assembled yet
            ret |= PROC_FN_RET_ABORT;
            return ret; // stop further processing
        }
    }

    ret |= ipProps->validityOK ? PROC_FN_RET_OK : PROC_FN_RET_ABORT;
    return ret;
}

static uint16_t nextIpIdentification = 0;

void insert_ipv4_header(uint8_t *hdr, const PcktHeaderElement *headers, EthInterface *intf) {
    uint8_t *hdrOrig = hdr;
    IPv4Props *ipProps = (IPv4Props *) &headers->props;
    ipProps->Identification = nextIpIdentification++; // auto-insert identification
    uint8_t version_length = (ipProps->Version << 4) | (ipProps->IHL & 0x0F);
    FILL_ADVANCE(hdr, &version_length, 1);
    FILL_ADVANCE(hdr, &ipProps->DSF, 1);
    FILL_WORD_H2N_ADVANCE(hdr, ipProps->TotalLength);
    FILL_WORD_H2N_ADVANCE(hdr, ipProps->Identification);
    uint16_t flags_fragOffset = ((ipProps->Flags & 0x07) << 13) | (ipProps->FragmentOffset & ~(0x07 << 13));
    FILL_WORD_H2N_ADVANCE(hdr, flags_fragOffset); // TODO...
    FILL_ADVANCE(hdr, &ipProps->TTL, 1);
    FILL_ADVANCE(hdr, &ipProps->Protocol, 1);
    uint8_t *ChkSumPtr = hdr;
    FILL_WORD_H2N_ADVANCE(hdr, 0x0000);
    FILL_ADVANCE(hdr, &ipProps->SourceIPAddr, 4);
    FILL_ADVANCE(hdr, &ipProps->DestIPAddr, 4);

    // calculate checksum after filling header if needed
    if (!(intf->capabilities & ETHINF_CAP_TX_IPCHKSUM_OFFLOAD)) {
        ipProps->HeaderChecksum = chksum(hdrOrig, ETH_IPv4_HEADER_SIZE, false);
        memcpy(ChkSumPtr, &ipProps->HeaderChecksum, 2);
    }
}

void ethmc_from_ipmc(uint8_t *hwa, ip4_addr ipa) {
    memcpy(hwa, ETH_ETHERNET_IPMC_HWADDR, ETH_HW_ADDR_LEN);
    const uint8_t *ipab = (const uint8_t *) &ipa;
    hwa[3] = (hwa[3] & 0x80) | (ipab[1] & 0x7F);
    hwa[4] = ipab[2];
    hwa[5] = ipab[3];
}

void ipv4_fill_props(IPv4Props *ipProps,
                     uint32_t payloadSize,
                     uint16_t protocol,
                     ip4_addr srcIpA,
                     ip4_addr dstIpA) {

    ipProps->IHL = ETH_IPv4_HEADER_SIZE / 4;
    ipProps->Version = 4;
    ipProps->DSF = 0x00;
    ipProps->TotalLength = ETH_IPv4_HEADER_SIZE + payloadSize;
    // Identification is filled on header insertion
    ipProps->Flags = 0x00;
    ipProps->FragmentOffset = 0x00;
    ipProps->TTL = 255;
    ipProps->Protocol = protocol;
    // HeaderChecksum is calculated on header insertion
    ipProps->SourceIPAddr = srcIpA;
    ipProps->DestIPAddr = dstIpA;

    // common fields for packet assembly
    ipProps->hdrInsFn = insert_ipv4_header;
    ipProps->headerSize = ETH_IPv4_HEADER_SIZE;

}

void ipv4_fill_in_chksum(uint8_t * hdr) {
    uint8_t * ChkSumPtr = hdr + 10; // create a pointer to checksum position
    uint16_t checksum = chksum(hdr, ETH_IPv4_HEADER_SIZE, false); // compute checksum
    memcpy(ChkSumPtr, &checksum, 2); // copy checksum to the proper position
}