#ifndef F20CF38F_7327_4608_8307_6AE058041CD5
#define F20CF38F_7327_4608_8307_6AE058041CD5

#include <algorithm>
#include <cstring>
#include <list>
#include <mesytec-mnode/mnode_cpp_types.h>
#include <mesytec-mnode/mnode_math.h>
#include <mesytec-mvlc/util/fmt.h>
#include <numeric>
#include <vector>

namespace mesytec::mnode::mana
{

class Arena
{
  public:
    static constexpr size_t initial_segment_size = 1u << 20;
    static constexpr size_t default_align_to = 8;
    static constexpr size_t default_max_segments = 1;

    Arena() = default;
    ~Arena() = default;

    Arena(Arena &&other) = default;
    Arena &operator=(Arena &&other) = default;

    // pushes at least the required amount of memory onto the arena.
    // 'required' is rounded up to the nearest multiple of 'default_align_to' to make
    // subsequent allocations aligned.
    u8 *push_size(size_t required)
    {
        size_t padded = round_up(required, default_align_to);
        size_t pad_waste = padded - required;
        required = padded;

        if (!segment_ || segment_->free() < required)
        {
            if (segment_count() < max_segments())
            {
                // could waste complete segments if required > default size of 1u << 20 but that's
                // fine for now
                segment_ = &segments_.emplace_back(
                    Arena::Segment(std::max(required, initial_segment_size)));
            }
            else
            {
                throw std::bad_alloc();
            }
        }

        assert(segment_ && segment_->free() >= required);

        auto result = segment_->cur();
        std::memset(result, 0, required);
        segment_->used += required;
        ++allocations_;
        pad_waste_ += pad_waste;
        return result;
    }

    template <typename T> T *push_t(size_t count = 1)
    {
        static_assert(std::is_trivial<T>::value, "T must be a trivial type");
        return reinterpret_cast<T *>(push_size(sizeof(T) * count));
    }

    char *push_cstr(const std::string &str)
    {
        auto mem = push_size(str.size() + 1);
        std::memcpy(mem, str.c_str(), str.size());
        return reinterpret_cast<char *>(mem);
    }

    void reset()
    {
        std::for_each(std::begin(segments_), std::end(segments_),
                      [](Segment &seg) { seg.reset(); });
        segment_ = segments_.empty() ? nullptr : &segments_.front();
        allocations_ = 0;
        pad_waste_ = 0;
    }

    size_t max_segments() const { return max_segments_; }
    size_t allocations() const { return allocations_; }
    size_t pad_waste() const { return pad_waste_; }
    size_t segment_count() const { return segments_.size(); }
    size_t capacity() const
    {
        auto accu = [](size_t sum, const Segment &seg) { return sum + seg.data.size(); };
        return std::accumulate(std::begin(segments_), std::end(segments_), static_cast<size_t>(0u),
                               accu);
    }

    size_t used() const
    {
        auto accu = [](size_t sum, const Segment &seg) { return sum + seg.used; };
        return std::accumulate(std::begin(segments_), std::end(segments_), static_cast<size_t>(0u),
                               accu) -
               pad_waste_;
    }

    const u8 *cur_begin() const { return segment_ ? segment_->data.data() : nullptr; }
    const u8 *cur_end() const { return segment_ ? segment_->cur() : nullptr; }

  private:
    Arena(const Arena &) = delete;
    Arena &operator=(const Arena &) = delete;

    struct Segment
    {
        std::vector<u8> data;
        size_t used;

        explicit Segment(size_t size)
            : data(size)
            , used(0)
        {
        }
        size_t free() const { return data.size() - used; }
        u8 *cur() { return data.data() + used; }
        void reset() { used = 0; }
    };

    size_t max_segments_ = default_max_segments;
    std::list<Segment> segments_;
    Segment *segment_ = nullptr;
    size_t allocations_ = 0;
    size_t pad_waste_ = 0;
};

inline std::string arena_stats(const Arena &arena)
{
    return fmt::format("allocations={}, capacity={}, pad_waste={}, used_size={}",
                       arena.allocations(), arena.capacity(), arena.pad_waste(), arena.used());
}

} // namespace mesytec::mnode::mana

#endif /* F20CF38F_7327_4608_8307_6AE058041CD5 */