#include <stddef.h>
#include <memory.h>

#include "ip_assembler.h"
#include "../../../dynmem.h"
#include "../../../global_state.h"

IPv4Assembler *ipra_new(EthInterface * intf) {
    IPv4Assembler *ipra = (IPv4Assembler *) dynmem_alloc(sizeof(IPv4Assembler));
    ipra->chains = NULL;
    ipra->intf = intf;
    return ipra;
}

static IPv4Fragment *ipra_all_frags_recvd(IPv4Assembler *ipra) {
    IPv4Fragment *fullChain = NULL;
    FragChainDesc *chainIter = ipra->chains;
    while ((chainIter != NULL) && (fullChain == NULL)) { // iter over all chains
        IPv4Fragment *fragIter = chainIter->fragChain;
        uint16_t expectedOffset = 0; // check that all fragments corresponding to the same IP packet have been received
        bool continuityOK = true; // we have possessed each fragments contiguously so far
        while ((fragIter != NULL) && continuityOK) { // iterate over all segments
            if (fragIter->offset == expectedOffset) {
                if ((fragIter->next == NULL) && (fragIter->lastFrag)) { // if last stored segment is not the chain's end, then the packet is not full
                    continuityOK = false;
                }
                expectedOffset += fragIter->size;
            } else {
                continuityOK = false;
            }
            fragIter = fragIter->next;
        }

        if (continuityOK) { // if we have a full chain, then return with the pointer to it
            fullChain = chainIter->fragChain;
        }
    }
    return fullChain;
}

static FragChainDesc *ipra_frag_chain_by_id(IPv4Assembler *ipra, uint16_t id) {
    FragChainDesc *fragChain = NULL, *iter = ipra->chains;
    while ((iter != NULL) && (fragChain == NULL)) {
        if (iter->id == id) {
            fragChain = iter;
        }
        iter = iter->next;
    }
    return fragChain;
}

static void ipra_remove_chain(IPv4Assembler * ipra, uint16_t id) {
    FragChainDesc * chainDesc = ipra_frag_chain_by_id(ipra, id);
    if (chainDesc == NULL) {
        return;
    }

    timer_unsched(E.tmr, chainDesc->id); // remove timeout schedule
    IPv4Fragment * iter = chainDesc->fragChain; // unlink and release fragment chain
    while (iter != NULL) {
        IPv4Fragment * oldIter = iter;
        iter = iter->next;
        dynmem_free(oldIter);
    }

    FragChainDesc * chainIter = ipra->chains;
    while (chainIter != NULL) { // unlink chain from chain list
        if (chainIter->next == chainDesc) { // if element is found
            chainIter->next = chainDesc->next; // skip by looked up element by setting next properly
            break; // break the loop
        }
        chainIter = chainIter->next;
    }

    // replace root element with the next one
    if (ipra->chains == chainDesc) {
        ipra->chains = chainDesc->next;
    }

    dynmem_free(chainDesc); // release chain descriptor
}

#define IP_REASSEMBLY_TIMEOUT_US (1000000)

static void ipra_timeout(Timer * tmr, AlarmUserData user) {
    IPv4Assembler * ipra = user.ptr;
    uint16_t id = user.u;
    ipra_remove_chain(ipra, id);
}

// insert into chain, without checking ID!
static void ipra_insert_into_chain(FragChainDesc *chainDesc, const IPv4Props *ipProps, const uint8_t *payload, uint32_t size) {
    IPv4Fragment *frag = (IPv4Fragment *) dynmem_alloc(sizeof(IPv4Fragment) + size); // allocate fragment
    frag->size = size;
    frag->lastFrag = !(ipProps->Flags & IP_FLAG_MF);
    frag->offset = ipProps->FragmentOffset;
    memcpy(frag->payload, payload, size);
    frag->next = NULL;

    IPv4Fragment *iter = chainDesc->fragChain;
    if (iter != NULL) { // if there are some segments present
        while (iter != NULL) {
            if (((iter->offset + iter->size) <= frag->offset) && // normal, sorted insertion
                ((iter->next == NULL) || (iter->next->offset >= (frag->offset + frag->size)))) {
                frag->next = iter->next;
                iter->next = frag;
            } else if ((iter == chainDesc->fragChain) && (frag->offset < iter->offset)) { // insert before the firstly received segment
                frag->next = iter;
                chainDesc->fragChain = frag;
            }
            iter = iter->next;
        }
    } else {
        chainDesc->fragChain = frag;
    }
    chainDesc->fullSize += size;
}

void ipra_input(IPv4Assembler *ipra, const IPv4Props *ipProps, const uint8_t *payload, uint32_t size) {
    uint16_t id = ipProps->Identification;
    FragChainDesc *chainDesc = ipra_frag_chain_by_id(ipra, id);
    if (chainDesc == NULL) { // if chain was found, allocate new one
        chainDesc = (FragChainDesc *) dynmem_alloc(sizeof(FragChainDesc));
        chainDesc->id = ipProps->Identification;
        chainDesc->alarmID = 0;
        chainDesc->fullSize = 0;
        chainDesc->fragChain = NULL;

        // link-in newly allocated fragment chain (position does not carry meaning, that's why insertion at the beginning is the simplest)
        chainDesc->next = ipra->chains;
        ipra->chains = chainDesc;
    }

    // insert segment into chain
    ipra_insert_into_chain(chainDesc, ipProps, payload, size);

    // reschedule timeout
    if (chainDesc->alarmID != 0) {
        timer_unsched(E.tmr, chainDesc->alarmID);
    }
    AlarmUserData aud;
    aud.ptr = ipra;
    aud.u = chainDesc->id;
    timer_sched_rel(E.tmr, IP_REASSEMBLY_TIMEOUT_US, ipra_timeout, aud);
    timer_report(E.tmr);
}

bool ipra_try_reassemble(IPv4Assembler *ipra, uint16_t id, uint8_t **payload, uint16_t *size, PcktHeaderElement *pcktHdrLe) {
    FragChainDesc *chainDesc = ipra_frag_chain_by_id(ipra, id);
    if (chainDesc == NULL) { // chain not found
        return false;
    }

    // check fragment continuity
    IPv4Fragment *fragIter = chainDesc->fragChain;
    uint16_t expectedOffset = 0; // check that all fragments corresponding to the same IP packet have been received
    bool continuityOK = true; // we have possessed each fragments contiguously so far
    while ((fragIter != NULL) && continuityOK) { // iterate over all segments
        if (fragIter->offset == expectedOffset) {
            if ((fragIter->next == NULL) && !(fragIter->lastFrag)) { // if last stored segment is not the chain's end, then the packet is not full
                continuityOK = false;
            }
            expectedOffset += fragIter->size;
        } else {
            continuityOK = false;
        }
        fragIter = fragIter->next;
    }

    if (continuityOK) { // if we have a full chain, then reassemble
        uint16_t offset = ETH_ETHERNET_HEADER_SIZE + ETH_IPv4_HEADER_SIZE; // offset for Ethernet and IPv4 header
        *payload = dynmem_alloc(chainDesc->fullSize + offset); // allocate space for the whole packet
        *size = chainDesc->fullSize;

        uint8_t * p = (*payload);
        IPv4Fragment * iter = chainDesc->fragChain;
        while (iter != NULL) { // assemble packet
            memcpy(p + iter->offset + offset, iter->payload, iter->size);
            iter = iter->next;
        }

        // insert headers and change fields to obfuscate defragmentation
        IPv4Props * ipProps = HEADER_FETCH_PROPS(IPv4Props, pcktHdrLe);
        ipProps->Flags = 0;
        ipProps->FragmentOffset = 0;
        ipProps->TotalLength = chainDesc->fullSize + ETH_IPv4_HEADER_SIZE;

        // release chain
        ipra_remove_chain(ipra, id);

        insert_ethernet_header(p, pcktHdrLe->prev, ipra->intf);
        insert_ipv4_header(p + ETH_ETHERNET_HEADER_SIZE, pcktHdrLe, ipra->intf);

        return true;
    }

    return false;
}