// SPDX-License-Identifier: GPL-2.0-only
/*
 * Digital Voice Modem - Common Library
 * GPLv2 Open Source. Use is subject to license terms.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 *  Copyright (C) 2025 Bryan Biedenkapp, N2PLL
 *
 */
#include "Defines.h"
#include "network/PacketBuffer.h"
#include "Log.h"
#include "Utils.h"

using namespace network;
using namespace compress;

#include <cassert>

// ---------------------------------------------------------------------------
//  Constants
// ---------------------------------------------------------------------------

#define MAX_FRAGMENT_SIZE 8192 * 1024 // 8MB max


// ---------------------------------------------------------------------------
//  Public Class Members
// ---------------------------------------------------------------------------

/* Initializes a new instance of the PacketBuffer class. */

PacketBuffer::PacketBuffer(bool compression, const char* name) :
    fragments(),
    m_compression(compression),
    m_name(name)
{
    /* stub */
}

/* Finalizes a instance of the PacketBuffer class. */

PacketBuffer::~PacketBuffer()
{
    clear();
}

/* Decode a network packet fragment. */

bool PacketBuffer::decode(const uint8_t* data, uint8_t** message, uint32_t* outLength)
{
    assert(data != nullptr);

    *message = nullptr;
    if (outLength != nullptr) {
        *outLength = 0;
    }

    uint8_t curBlock = data[8U];
    uint8_t blockCnt = data[9U];

    Fragment* frag = new Fragment();
    frag->compressedSize = 0U;
    frag->size = 0U;
    frag->blockSize = 0U;
    frag->blockId = 0U;
    frag->data = nullptr;

    // if this is the first block store sizes and initialize temp buffer
    if (curBlock == 0U) {
        uint32_t size = GET_UINT32(data, 0U);
        uint32_t compressedSize = GET_UINT32(data, 4U);

        // make sure we can't exceed max fragment size -- prevent potential DOS attack by sending
        // enormous fragment sizes
        if (size > MAX_FRAGMENT_SIZE || compressedSize > MAX_FRAGMENT_SIZE) {
            LogError(LOG_NET, "%s, fragment size exceeds maximum. BUGBUG.", m_name);
            delete frag;
            return false;
        }

        frag->size = size;
        frag->compressedSize = compressedSize;
    }

    // scope is intentional
    {
        frag->blockId = curBlock;
        frag->data = new uint8_t[FRAG_BLOCK_SIZE + 1U];

        ::memcpy(frag->data, data + FRAG_HDR_SIZE, FRAG_BLOCK_SIZE);
        // Utils::dump(1U, "PacketBuffer::decode(), Block Payload", frag->data, FRAG_BLOCK_SIZE);
        fragments.insert(curBlock, frag);
    }

    LogInfoEx(LOG_NET, "%s, Inbound Packet Fragment, block %u of %u, rxFragments = %u", m_name, curBlock, blockCnt, fragments.size());

    // do we have all the blocks?
    if (fragments.size() == blockCnt + 1U) {
        fragments.lock(false);
        auto firstIt = fragments.find(0U);
        if (firstIt == fragments.end() || firstIt->second == nullptr) {
            LogError(LOG_NET, "%s, Packet Fragment, error missing block 0? Packet dropped.", m_name);
            fragments.unlock();
            clear();
            return false;
        }

        Fragment* firstFrag = firstIt->second;

        if (firstFrag->size == 0U) {
            LogError(LOG_NET, "%s, Packet Fragment, error missing size information", m_name);
            fragments.unlock();
            clear();
            return false;
        }

        if (firstFrag->compressedSize == 0U) {
            LogError(LOG_NET, "%s, Packet Fragment, error missing compressed size information", m_name);
            fragments.unlock();
            clear();
            return false;
        }

        uint32_t compressedLen = firstFrag->compressedSize;
        uint32_t len = firstFrag->size;

        if (compressedLen > MAX_FRAGMENT_SIZE || len > MAX_FRAGMENT_SIZE) {
            LogError(LOG_NET, "%s, Packet Fragment, error invalid packet length metadata", m_name);
            fragments.unlock();
            clear();
            return false;
        }

        // Reassembly buffer must be compressed length, not uncompressed length.
        DECLARE_UINT8_ARRAY(buffer, compressedLen + 1U);
        ::memset(buffer, 0x00U, compressedLen + 1U);

        if (fragments.size() == 1U) {
            uint32_t copyLen = (compressedLen > FRAG_BLOCK_SIZE) ? FRAG_BLOCK_SIZE : compressedLen;
            ::memcpy(buffer, firstFrag->data, copyLen);
        } else {
            for (uint8_t i = 0U; i <= blockCnt; i++) {
                auto it = fragments.find(i);
                if (it == fragments.end() || it->second == nullptr || it->second->data == nullptr) {
                    LogError(LOG_NET, "%s, Packet Fragment, error missing block %u, packet dropped", m_name, i);
                    fragments.unlock();
                    clear();
                    return false;
                }

                uint32_t offs = i * FRAG_BLOCK_SIZE;
                if (offs >= compressedLen)
                    break;

                uint32_t copyLen = FRAG_BLOCK_SIZE;
                if (offs + FRAG_BLOCK_SIZE > compressedLen)
                    copyLen = compressedLen - offs;

                ::memcpy(buffer + offs, it->second->data, copyLen);
            }
        }

        if (m_compression) {
            uint32_t decompressedLen = 0U;
            UInt8Array decompressed = Compression::decompress(buffer, compressedLen, &decompressedLen);
            *message = new uint8_t[decompressedLen];
            ::memset(*message, 0x00U, decompressedLen);
            ::memcpy(*message, decompressed.get(), decompressedLen);

            // check that we got the appropriate data
            if (decompressedLen == len && message != nullptr) {
                if (outLength != nullptr) {
                    *outLength = len;
                }

                fragments.unlock();
                clear();
                return true;
            }
        }
        else {
            uint8_t* outBuf = new uint8_t[len + 1U];
            ::memset(outBuf, 0x00U, len + 1U);
            ::memcpy(outBuf, buffer, len);

            *message = outBuf;
            if (outLength != nullptr) {
                *outLength = len;
            }

            fragments.unlock();
            clear();
            return true;
        }

        fragments.unlock();
    }

    return false;
}

/* Encode a network packet fragment. */

void PacketBuffer::encode(uint8_t* data, uint32_t length)
{
    assert(data != nullptr);
    assert(length > 0U);

    // erase any buffered fragments
    clear();

    // create temporary buffer
    uint32_t compressedLen = 0U;
    UInt8Array buffer = nullptr;
    if (m_compression) {
        buffer = Compression::compress(data, length, &compressedLen);
    } else {
        buffer = std::make_unique<uint8_t[]>(length + 1U);
        ::memset(buffer.get(), 0x00U, length + 1U);
        ::memcpy(buffer.get(), data, length);
        compressedLen = length;
    }

    // create packet fragments
    uint8_t blockCnt = (compressedLen / FRAG_BLOCK_SIZE) + (compressedLen % FRAG_BLOCK_SIZE ? 1U : 0U);
    uint32_t offs = 0U;
    for (uint8_t i = 0U; i < blockCnt; i++) {
        // build dataset
        uint16_t bufSize = FRAG_SIZE;

        Fragment* frag = new Fragment();
        frag->blockId = i;

        frag->data = new uint8_t[bufSize];
        ::memset(frag->data, 0x00U, bufSize);

        if (i == 0U) {
            SET_UINT32(length, frag->data, 0U);
            frag->size = length;
            SET_UINT32(compressedLen, frag->data, 4U);
            frag->compressedSize = compressedLen;
        }

        frag->data[8U] = i;
        frag->data[9U] = blockCnt - 1U;

        uint32_t blockSize = FRAG_BLOCK_SIZE;
        if (offs + FRAG_BLOCK_SIZE > compressedLen)
            blockSize = FRAG_BLOCK_SIZE - ((offs + FRAG_BLOCK_SIZE) - compressedLen);

        ::memcpy(frag->data + FRAG_HDR_SIZE, buffer.get() + offs, blockSize);

        offs += FRAG_BLOCK_SIZE;

        fragments.insert(i, frag);
        LogInfoEx(LOG_NET, "%s, Outbound Packet Fragment, block %u of %u, txFragments = %u", m_name, i, blockCnt - 1U, fragments.size());
    }
}

/* Helper to clear currently buffered fragments. */

void PacketBuffer::clear()
{
    fragments.lock(false);
    for (auto entry : fragments) {
        Fragment* frag = entry.second;
        if (frag != nullptr) {
            if (frag->data != nullptr)
                delete[] frag->data;
            frag->data = nullptr;
            delete frag;
        }
    }
    fragments.unlock();

    fragments.clear();
}
