//
// Created by epagris on 2023.01.14..
//

#include "tcp_segment.h"
#include "../../utils.h"
#include "ethernet_frame.h"
#include "../../dynmem.h"
#include "../../eth_interface.h"

#include "tcp_udp_common.h"
#include "ipv4_packet.h"

#define TCP_SINGLE_BYTE_OPTION(kind) (((kind) == TCP_OPT_KIND_EOL) || ((kind) == TCP_OPT_KIND_NOP))

static int tcp_fetch_options(const uint8_t * hdr, uint8_t size, TcpOption ** optLe) {
    uint8_t idx = 0;
    TcpOption * lastOpt = NULL;
    int optCnt = 0;
    while (idx < size) {
        uint8_t kind = hdr[idx + 0]; // read option kind

        // skip further reading when EOL or NOP has been detected
        if (TCP_SINGLE_BYTE_OPTION(kind)) {
            idx++;
            continue;
        }

        uint8_t valueSize = hdr[idx + 1] - 2; // read option length

        TcpOption * opt = (TcpOption *) dynmem_alloc(sizeof(TcpOption) + valueSize); // allocate option

        // fill-in
        opt->kind = kind;
        opt->size = valueSize;
        memcpy(opt->value, hdr + idx + 2, opt->size); // copy option value
        opt->next = NULL;

        // store option
        if (*optLe == NULL) {
            *optLe = opt;
        } else {
            lastOpt->next = opt;
        }

        lastOpt = opt; // advance last opt
        idx += opt->size + 2; // advance index
        optCnt++; // increase opt count
    }

    return optCnt;
}

static void tcp_insert_options(uint8_t * hdr, uint8_t size, TcpOption * optLe) {
    TcpOption * opt = optLe;
    uint8_t idx = 0;
    while (opt != NULL && idx < size) {
        hdr[idx + 0] = opt->kind; // insert kind

        // skip further reading when EOL or NOP has been detected
        if (TCP_SINGLE_BYTE_OPTION(opt->kind)) {
            idx++;
            opt = opt->next;
            continue;
        }

        uint8_t length = opt->size + 2; // insert length
        hdr[idx + 1] = length;
        memcpy(hdr + idx + 2, opt->value, opt->size); // insert value

        idx += length; // advance idx
        opt = opt->next; // advance iterator
    }

    // insert NOP padding if needed
    if (idx < (size - 1)) {
        uint8_t paddingNop = size - idx - 1;
        memset(hdr + idx, 1, paddingNop);
    }

    // insert EOL if needed
    if (idx < size) {
        memset(hdr + idx, 0, 1);
    }
}

uint32_t tcp_get_options_size(TcpOption * optLe) {
    TcpOption * opt = optLe;
    uint32_t size = 0;
    while (opt != NULL) {
        if (TCP_SINGLE_BYTE_OPTION(opt->kind)) {
            size += 1;
        } else {
            size += opt->size + 2;
        }
        opt = opt->next;
    }

    // round up to the neareast integer multiple of four bytes
    size = CEIL_TO_4(size);

    return size;
}

TcpOption * tcp_get_option_by_kind(TcpOption * opt, uint16_t kind) {
    TcpOption * res = NULL, * iter = opt;
    while (iter != NULL) {
        if (iter->kind == kind) {
            res = opt;
        }
        iter = iter->next;
    }
    return res;
}

int parse_tcp(const uint8_t *hdr, uint32_t size, PcktHeaderElement *pcktHdrLe, struct EthInterface_ *intf, PcktProcFnPassbackData *pb) {
    const uint8_t * hdrBegin = hdr;

    // fetch fixed portion of the segment header
    TcpProps * tcpProps = HEADER_FETCH_PROPS(TcpProps, pcktHdrLe);
    FETCH_WORD_H2N_ADVANCE(&tcpProps->SourcePort, hdr);
    FETCH_WORD_H2N_ADVANCE(&tcpProps->DestinationPort, hdr);
    FETCH_DWORD_H2N_ADVANCE(&tcpProps->SequenceNumber, hdr);
    FETCH_DWORD_H2N_ADVANCE(&tcpProps->AcknowledgementNumber, hdr);
    uint8_t dataOffset_Flags;
    FETCH_BYTE_ADVANCE(&dataOffset_Flags, hdr);
    tcpProps->DataOffset = (dataOffset_Flags >> 4) << 2; // eliminate non-relevant bits
    uint8_t flagsLower;
    FETCH_BYTE_ADVANCE(&flagsLower, hdr);
    tcpProps->Flags = flagsLower | ((dataOffset_Flags & 0x01) << 8);
    FETCH_WORD_H2N_ADVANCE(&tcpProps->Window, hdr);
    FETCH_WORD_H2N_ADVANCE(&tcpProps->Checksum, hdr);
    FETCH_WORD_H2N_ADVANCE(&tcpProps->UrgentPtr, hdr);

    // fetch options
    tcpProps->options = NULL; // no options by default
    uint8_t fullHeaderSize = tcpProps->DataOffset;
    if (fullHeaderSize > ETH_TCP_HEADER_SIZE) {
        uint8_t optSize = fullHeaderSize - ETH_TCP_HEADER_SIZE;
        tcp_fetch_options(hdr, optSize, &tcpProps->options);
    }

    // verify checksum
    if (!(intf->capabilities & ETHINF_CAP_RX_TCPUDPCHKSUM_OFFLOAD)) { // TODO
        const IPv4Props *ipProps = (IPv4Props *) &pcktHdrLe->prev->props;
        uint16_t headerAndPayloadLength = ipProps->bytesToEnd;
        IPv4PseudoHeader ph = {ipProps->SourceIPAddr, ipProps->DestIPAddr, 0, ETH_TCP_PACKET_CLASS, htons(headerAndPayloadLength)};
        //uint16_t chkSum = tcp_udp_checksum(&ph, hdrBegin, headerAndPayloadLength);
        uint16_t chkSum = 0; // TODO ....
        tcpProps->validityOK = (chkSum == 0);
    } else {
        tcpProps->validityOK = true;
    }

    tcpProps->headerSize = fullHeaderSize;
    tcpProps->containedPacketClass = 0;

    return tcpProps->validityOK ? PROC_FN_RET_OK : PROC_FN_RET_ABORT;
}

void insert_tcp_header(uint8_t *hdr, const PcktHeaderElement *headers, EthInterface *intf) {
    uint8_t * hdrBegin = hdr;
    TcpProps * tcpProps = HEADER_FETCH_PROPS(TcpProps, headers); // fetch header

    // get option area size and computer full area size
    uint32_t optSize = tcp_get_options_size(tcpProps->options);
    uint32_t fullHeaderSize = ETH_TCP_HEADER_SIZE + optSize;
    tcpProps->DataOffset = fullHeaderSize >> 2;

    // fill beginning of the header
    FILL_WORD_H2N_ADVANCE(hdr, tcpProps->SourcePort);
    FILL_WORD_H2N_ADVANCE(hdr, tcpProps->DestinationPort);
    FILL_DWORD_H2N_ADVANCE(hdr, tcpProps->SequenceNumber);
    FILL_DWORD_H2N_ADVANCE(hdr, tcpProps->AcknowledgementNumber);
    uint8_t dataOffset_Flags = (tcpProps->DataOffset << 4) | (tcpProps->Flags >> 8);
    FILL_BYTE_ADVANCE(hdr, &dataOffset_Flags);
    uint8_t flagsLower = tcpProps->Flags & 0xFF;
    FILL_BYTE_ADVANCE(hdr, &flagsLower);
    FILL_WORD_H2N_ADVANCE(hdr, tcpProps->Window);
    tcpProps->Checksum = 0;
    uint8_t * ChkSumPtr = hdr;
    FILL_WORD_H2N_ADVANCE(hdr, tcpProps->Checksum);
    FILL_WORD_H2N_ADVANCE(hdr, tcpProps->UrgentPtr);

    // fill options
    tcp_insert_options(hdr, optSize, tcpProps->options);

    // calculate checksum
    if (!(intf->capabilities & ETHINF_CAP_TX_TCPUDPCHKSUM_OFFLOAD)) {
        const IPv4Props *ipProps = (IPv4Props *) &headers->prev->props;
        uint16_t headerAndPayloadLength = fullHeaderSize + headers->props.bytesToEnd;
        IPv4PseudoHeader ph = {ipProps->SourceIPAddr, ipProps->DestIPAddr, 0, ETH_TCP_PACKET_CLASS, htons(headerAndPayloadLength)};
        tcpProps->Checksum = tcp_udp_checksum(&ph, hdrBegin, headerAndPayloadLength);
        memcpy(ChkSumPtr, &tcpProps->Checksum, 2);
    }
}

// ----------------------------------

TcpOption *tcpopt_new(uint32_t kind, uint32_t size) {
    TcpOption * opt = dynmem_alloc(sizeof(TcpOption) + size);
    opt->size = size;
    opt->kind = kind;
    opt->next = NULL;
    return opt;
}