#include "IpPacketFragmentation.h"
#include <algorithm>
#include <tins/ipv6.h>
#include <tins/rawpdu.h>
#include <tins/tcp.h>

const size_t IpPacketFragmentation::fragmentionHeadersize = sizeof (FragmentionHeaderUnion) - 1;

IpPacketFragmentation::IpPacketFragmentation(const size_t newMtu) : mtu(newMtu), idCounter()
{
}


IpPacketFragmentation::~IpPacketFragmentation()
{

}

void IpPacketFragmentation::addExtensionHeader(IN const uint8_t nextHeader, IN const uint8_t *startPtr, IN Tins::IPv6 & ipFragmentPdu, const size_t headersize)
{
    Tins::IPv6::ext_header extensionHeader(nextHeader, 6 , startPtr);
    ipFragmentPdu.add_ext_header(extensionHeader);
}

void IpPacketFragmentation::initFragmentationHeader(FragmentionHeaderStruct* ptrFragmentionHeaderStruct)
{
    ptrFragmentionHeaderStruct->NextHeader = Tins::Constants::IP::PROTO_TCP;
    ptrFragmentionHeaderStruct->Reserved = 0;
    ptrFragmentionHeaderStruct->Res = 0;
    ptrFragmentionHeaderStruct->MFlag = 1;
    ptrFragmentionHeaderStruct->Identification = ++idCounter;
}

bool IpPacketFragmentation::handle(IN const Tins::PDU & pdu, IN IPacketHandler *  callBackHandler)
{
    if (callBackHandler == nullptr || callBackHandler == this)
    {
        return false;
    }

    const size_t originPduSize = pdu.size();
    const size_t fragmentationCount = (originPduSize / mtu) + 1;
    if (fragmentationCount < 2)
    {
        return callBackHandler->handle(pdu, this);
    }

    Tins::IPv6 * ipPdu = pdu.clone()->find_pdu<Tins::IPv6>();
    if (ipPdu == nullptr)
    {
        return false;
    }

    Tins::PDU * ipDataPdu = ipPdu->inner_pdu();
    if (ipDataPdu == nullptr)
    {
        return false;
    }

    ByteVector ipPayload2 = ipPdu->serialize();
    ByteVector ipPayload(ipPayload2.begin() + ipPdu->header_size(), ipPayload2.end());
    uint8_t originNextHeader = ipPdu->next_header();
    const long ipPayloadSize = static_cast<const long>(ipPayload.size());
    const long fragmentationSize64B = (ipPayloadSize/ static_cast<const long>(fragmentationCount)) / 64;
    const long fragmentationSize = fragmentationSize64B * 64;
    const long payLoadEndIndicator = ipPayloadSize - fragmentationSize;
    ByteVector::iterator fragmentPosIt = ipPayload.begin();
    FragmentionHeaderUnion fragmentionHeaderUnion;
    FragmentionHeaderStruct * ptrFragmentionHeaderStruct = &fragmentionHeaderUnion.Structed;
    uint8_t * ptrStartFragmentionHeader = &fragmentionHeaderUnion.Bytes[2];
    initFragmentationHeader(ptrFragmentionHeaderStruct);

    long frameOffset;
    for (frameOffset = 0; frameOffset < payLoadEndIndicator; frameOffset += fragmentationSize)
    {
        ByteVector::iterator fragmentStart = fragmentPosIt;
        fragmentPosIt = fragmentStart + fragmentationSize;
        ptrFragmentionHeaderStruct->FragmentOffset = static_cast<uint16_t>(frameOffset) / 8;
        std::swap(fragmentionHeaderUnion.Bytes[2],fragmentionHeaderUnion.Bytes[3]);
        if (!createAndForwardFragmend(originNextHeader,pdu, fragmentStart, fragmentPosIt, ptrFragmentionHeaderStruct, ptrStartFragmentionHeader, callBackHandler))
        {
            return false;
        }

       std::swap(fragmentionHeaderUnion.Bytes[2],fragmentionHeaderUnion.Bytes[3]);
    }

    ptrFragmentionHeaderStruct->MFlag = 0;
    ptrFragmentionHeaderStruct->FragmentOffset = static_cast<uint16_t>(frameOffset) / 8;
    std::swap(fragmentionHeaderUnion.Bytes[2],fragmentionHeaderUnion.Bytes[3]);
    return createAndForwardFragmend(originNextHeader, pdu, fragmentPosIt, ipPayload.end(), ptrFragmentionHeaderStruct, ptrStartFragmentionHeader, callBackHandler);
}

bool IpPacketFragmentation::createAndForwardFragmend(IN const uint8_t originNextHeader, IN const Tins::PDU & pdu, IN const ByteVector::iterator & fragmentStart, IN const ByteVector::iterator & fragmentPosIt, IN FragmentionHeaderStruct * ptrFragmentionHeaderStruct, IN uint8_t * ptrStartFragmentionHeader, IN IPacketHandler *  callBackHandler)
{
    SPtrByteVector fragmentPayload = std::make_shared<ByteVector>(fragmentStart, fragmentPosIt);
    Tins::PDU * fragmentPdu = pdu.clone();
    if (fragmentPdu == nullptr)
    {
        return false;
    }

    Tins::IPv6 * ipFragmentPdu = fragmentPdu->find_pdu<Tins::IPv6>();
    if (ipFragmentPdu == nullptr)
    {
        return false;
    }

    addExtensionHeader(originNextHeader, ptrStartFragmentionHeader, *ipFragmentPdu, IpPacketFragmentation::fragmentionHeadersize);
    SPtrRawPDU rawFragmentPdu = std::make_shared<Tins::RawPDU>(fragmentPayload->data(), static_cast<uint32_t>(fragmentPayload->size()));
    ipFragmentPdu->inner_pdu(rawFragmentPdu.get());
    ipFragmentPdu->next_header(Tins::IPv6::FRAGMENT);
    return callBackHandler->handle(*fragmentPdu, this);
}