FeatherMC/src/NetworkMessage.h

#pragma once

#include "util/ByteUtil.h"
#include "Types.h"

#include <cstdint>
#include <cstring>
#include <vector>
#include <string>
#include <cassert>

#define VARINT_MAX_SIZE 5

namespace Feather
{
    class NetworkMessage
    {
    public:
        NetworkMessage(int32_t expectedSize)
        {
            m_data.reserve(expectedSize);
            m_data.resize(VARINT_MAX_SIZE);
        }

        template <typename T>
        inline void WriteData(const T* data, size_t size)
        {
            size_t offset = m_data.size();

            m_data.resize(offset + size);
            std::memcpy(&m_data[offset], data, size);
        }

        // Writes a value, converting to big endian by default
        template <typename T, Endian asEndian = Endian::Big>
        inline void Write(const T& value)
        {
            if constexpr (asEndian == Endian::Big && sizeof(T) > 1)
            {
                T valueReversed = ReverseBytes(value);
                WriteData(&valueReversed, sizeof(T));
            }
            else
            {
                WriteData(&value, sizeof(T));
            }
        }

        // Write bool as 1 byte
        template <>
        inline void Write<bool>(const bool& value)
        {
            Write<uint8_t>(value ? 0x01 : 0x00);
        }

        template <typename Int = int32_t>
        inline void WriteVarInt(Int val)
        {
            int32_t value = static_cast<int32_t>(val);
            do {
                uint8_t temp = uint8_t(value & 0b01111111);
                value >>= 7;
                if (value != 0)
                    temp |= 0b10000000;

                Write(temp);
            } while (value != 0);
        }

        inline void WriteString(const char* string, int32_t length)
        {
            int32_t stringSize = length;
            WriteVarInt(stringSize);
            WriteData(string, stringSize);
        }

        // Write string using std::string reference
        inline void WriteString(const std::string& string)
        {
            WriteString(string.c_str(), string.length());
        }

        // Write string using string_view
        inline void WriteString(std::string_view view)
        {
            WriteString(view.data(), view.size());
        }

        inline void WritePosition(BlockPos pos)
        {
            Write<int64_t>(pos.Encode());
        }

        inline void WriteSubMessage(const NetworkMessage& message)
        {
            WriteData(message.GetData(), message.GetDataSize());
        }

        inline void Finalize()
        {
            PrependVarIntSize();
#ifdef _DEBUG
            m_finalized = true;
#endif
        }

        inline const uint8_t* GetData() const
        {
#ifdef _DEBUG
            // You must call Finalize() first
            Assert(m_finalized, "NetworkMessage sent or resolved before Finalize() was called.");
#endif
            return &m_data[m_startOffset];
        }

        inline size_t GetDataSize() const
        {
#ifdef _DEBUG
            // You must call Finalize() first
            Assert(m_finalized, "NetworkMessage sent or resolved before Finalize() was called.");
#endif
            return m_data.size() - m_startOffset;
        }

    private:
        std::vector<uint8_t> m_data;
        size_t m_startOffset = 0;
#ifdef _DEBUG
        bool m_finalized = false;
#endif

        void PrependVarIntSize()
        {
            uint8_t sizeBytes[VARINT_MAX_SIZE];

            // compute the size of the packet excluding the size marker
            int32_t value = static_cast<int32_t>(m_data.size() - VARINT_MAX_SIZE);
            size_t index = 0;

            // perform WriteVarInt into the static buffer sizeBytes
            do {
                uint8_t temp = uint8_t(value & 0b01111111);
                value >>= 7;
                if (value != 0)
                    temp |= 0b10000000;

                sizeBytes[index++] = temp;
            } while (value != 0 && index < VARINT_MAX_SIZE);

            // where to start writing the size marker
            // if we used all 5 bytes, this is zero
            // if we used 3 bytes, this would be 2
            size_t startIndex = VARINT_MAX_SIZE - index;

            // Copy sizeBytes into our actual buffer, but align it to the right
            for (size_t i = 0; i < index; i++)
                m_data[startIndex + i] = sizeBytes[i];

            m_startOffset = startIndex;
        }
    };
}