Towards/modules/serialization/include/tw/serial/WorldStateWriter.hpp

#pragma once


#include "Codec.hpp"
#include "GlmCodec.hpp"
#include "EnttCodec.hpp"

#include <entt/entt.hpp>
#include <glm/vec3.hpp>
#include <cstdint>
#include <span>
#include <spdlog/spdlog.h>

namespace tw::serial {

class WorldStateWriter {
    BinaryWriter  m_w;

    // Offsets for length-prefix patching
    std::size_t m_entity_count_offset{ 0 };
    uint32_t    m_entity_count{ 0 };

public:
    explicit WorldStateWriter(BinaryBuffer& buf) noexcept : m_w(buf) {}

    void begin(uint32_t frame_idx, uint32_t tick_idx, uint32_t message_type) noexcept {
        m_entity_count = 0;

        // message_type — lets the receiver dispatch without peeking further
        m_w.encode<uint32_t>(message_type);

        // sequence — this message is never a reply to a request
        m_w.encode<uint32_t>(0);

        // frame_idx
        m_w.encode<uint32_t>(frame_idx);

        // tick_idx — counts the states that went out, so the receiver can tell
        // a late one from a new one. Says nothing about the frame answered.
        m_w.encode<uint32_t>(tick_idx);

        // entity_count placeholder — patched when end() is called
        m_entity_count_offset = m_w.reserve_u32();
    }

    template<typename Range>
    void write_spawns(const Range& spawns) noexcept {
        auto count = static_cast<uint32_t>(std::size(spawns));
        m_w.encode<uint32_t>(count);
        for (const entt::entity e : spawns) {
            m_w.encode<entt::entity>(e);
        }
    }

    template<typename Range>
    void write_despawns(const Range& despawns) noexcept {
        auto count = static_cast<uint32_t>(std::size(despawns));
        // m_w.encode<uint32_t>(count);
        for (const entt::entity e : despawns) {
            m_w.encode<entt::entity>(e);
        }
    }

    void write_entity(uint32_t id, const glm::vec3& pos) noexcept {
        m_w.write(id);
        // Write x, y, z as 3 contiguous floats
        m_w.write_bytes(&pos.x, 3 * sizeof(float));
        ++m_entity_count;
    }

    void write_entity(entt::entity entity, const glm::vec3& pos) noexcept {
        write_entity(static_cast<uint32_t>(entity), pos);
    }

    void end() noexcept {
        m_w.patch_u32(m_entity_count_offset, m_entity_count);
    }

    std::span<const std::byte> view() noexcept {
        return m_w.buffer().view();
    }

    void reset() noexcept {
        m_w.reset();
        m_entity_count = 0;
    }
};

struct WorldStateHeader {
    uint32_t packet_type;
    uint32_t frame_idx;
    uint32_t tick_idx;
    uint32_t entity_count;
};

struct EntityRecord {
    uint32_t  id;
    glm::vec3 position;
};

class WorldStateReader {
    BinaryReader m_r;
    WorldStateHeader m_header{};

    uint32_t m_spawn_count{ 0 };
    uint32_t m_spawns_read{ 0 };

    uint32_t m_despawn_count{ 0 };
    uint32_t m_despawns_read{ 0 };

    uint32_t m_entities_read{ 0 };

public:
    explicit WorldStateReader(std::span<const std::byte> data) noexcept
        : m_r(data) {}

    /** Read the header. Must be called first. */
    WorldStateHeader read_header() noexcept {
        // m_header.packet_type  = m_r.decode<uint32_t>();
        m_header.frame_idx    = m_r.decode<uint32_t>();
        m_header.tick_idx     = m_r.decode<uint32_t>();
        m_header.entity_count = m_r.decode<uint32_t>();

        // spawn count follows immediately
        m_spawn_count  = m_r.decode<uint32_t>();
        return m_header;
    }

    /** Read the next spawn entity id.  Returns 0 when exhausted. */
    bool has_spawn() const noexcept { return m_spawns_read < m_spawn_count; }

    uint32_t read_spawn() noexcept {
        assert(has_spawn());
        ++m_spawns_read;
        uint32_t id = m_r.decode<uint32_t>();
        // if (!has_spawn()) {
        //     // transition to despawns
        //     m_despawn_count = m_r.decode<uint32_t>();
        //     m_phase = Phase::Despawns;
        // }
        return id;
    }

    /** Skip remaining spawns and enter despawn phase. */
    void skip_spawns() noexcept {
        while (has_spawn()) read_spawn();
    }

    bool has_despawn() const noexcept { return m_despawns_read < m_despawn_count; }

    uint32_t read_despawn() noexcept {
        assert(has_despawn());
        ++m_despawns_read;
        uint32_t id = m_r.decode<uint32_t>();
        return id;
    }

    void skip_despawns() noexcept {
        while (has_despawn()) read_despawn();
    }

    bool has_entity() const noexcept {
        return m_entities_read < m_header.entity_count;
    }

    EntityRecord read_entity() noexcept {
        assert(has_entity());
        EntityRecord rec;
        rec.id = m_r.read<uint32_t>();
        m_r.read_bytes(&rec.position.x, 3 * sizeof(float));
        ++m_entities_read;
        return rec;
    }
};

} // namespace tw::serial