#pragma once

#include <cstdint>
#include <cstring>
#include <vector>

namespace Feather::Network
{
    class NetworkMessageCounter
    {
    public:

        inline void WriteVarInt(int32_t value)
        {
            do {
                uint8_t temp = uint8_t(value & 0b01111111);
                value >>= 7;
                if (value != 0)
                    temp |= 0b10000000;

                m_size++;
            } while (value != 0);
        }

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

            m_size += stringSize;
        }

        inline int32_t GetSize() const { return m_size; }

    private:
        int32_t m_size = 0;
    };

    class NetworkMessage
    {
    public:

        NetworkMessage(int32_t size)
        {
            NetworkMessageCounter packetSizeCounter;
            packetSizeCounter.WriteVarInt(size);

            m_data.reserve(size_t(size) + size_t(packetSizeCounter.GetSize()));
            WriteVarInt(size);
        }

        inline void WriteVarInt(int32_t value)
        {
            do {
                uint8_t temp = uint8_t(value & 0b01111111);
                value >>= 7;
                if (value != 0)
                    temp |= 0b10000000;

                m_data.push_back(temp);
            } while (value != 0);
        }

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

            size_t oldSize = m_data.size();
            size_t newSize = oldSize + size_t(stringSize);
            m_data.resize(newSize);
            std::memcpy(&m_data[oldSize], string, size_t(stringSize));
        }

        const uint8_t* GetData()     const { return m_data.data(); }
        size_t         GetDataSize() const { return m_data.size(); }

    private:
        std::vector<uint8_t> m_data;
    };

    #define DEFINE_MESSAGE_WRAPPER(x)                  \
        template <typename... Args>                    \
        NetworkMessage x(Args... args)                 \
        {                                              \
            NetworkMessageCounter counter;             \
            x(counter, args...);                       \
            NetworkMessage message(counter.GetSize()); \
            x(message, args...);                       \
            return message;                            \
        }
}