#pragma once

#if !defined(MIJIN_MEMORY_DATA_POOL_HPP_INCLUDED)
#define MIJIN_MEMORY_DATA_POOL_HPP_INCLUDED 1

#include <array>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <vector>

namespace mijin
{

//
// public defines
//

//
// public constants
//

//
// public types
//

template<std::size_t PAGE_SIZE = 4096ul * 1024ul>
struct DataPool
{
    static_assert(PAGE_SIZE % alignof(std::max_align_t) == 0);
    
    struct alignas(std::max_align_t) Page : std::array<std::uint8_t, PAGE_SIZE> {};

    using page_ptr_t = std::unique_ptr<Page>;

    std::vector<page_ptr_t> pages;
    std::size_t offset = 0;

    DataPool() = default;
    DataPool(const DataPool&) = delete;
    DataPool(DataPool&&) noexcept = default;

    DataPool& operator=(const DataPool&) = delete;
    DataPool& operator=(DataPool&&) noexcept = default;

    void* allocate(std::size_t bytes, std::size_t alignment = 1);
    inline void reset() { offset = 0; }

    template<typename T>
    T* allocateAs(std::size_t elements = 1) {
        return static_cast<T*>(allocate(sizeof(T) * elements, alignof(T)));
    }
};

//
// public functions
//

template<std::size_t PAGE_SIZE>
void* DataPool<PAGE_SIZE>::allocate(std::size_t bytes, std::size_t alignment)
{
    assert(bytes > 0 && bytes <= PAGE_SIZE);
    assert(alignment > 0 && alignment <= alignof(std::max_align_t));

    if (offset % alignment != 0) {
        offset += alignment - (offset % alignment);
    }

    const std::size_t remainingOnPage = PAGE_SIZE - (offset % PAGE_SIZE);
    const std::size_t page = offset / PAGE_SIZE;
    const std::size_t localOffset = offset % PAGE_SIZE;
    if (remainingOnPage == PAGE_SIZE || remainingOnPage < bytes)
    {
        // next page
        if (page + 1 >= pages.size()) {
            pages.push_back(std::make_unique<Page>());
        }
        offset = PAGE_SIZE * (pages.size() - 1);
    }
    std::uint8_t* result = &(*pages[page])[localOffset];
    offset += bytes;
    assert(reinterpret_cast<std::uintptr_t>(result) % alignment == 0);
    return result;
}

} // namespace mijin

#endif // !defined(MIJIN_MEMORY_DATA_POOL_HPP_INCLUDED)