#pragma once

#if !defined(MIJIN_MEMORY_STACK_ALLOCATOR_HPP_INCLUDED)
#define MIJIN_MEMORY_STACK_ALLOCATOR_HPP_INCLUDED 1

#include "../debug/assert.hpp"
#include "../internal/common.hpp"
#include "../util/align.hpp"
#include "../util/traits.hpp"

namespace mijin
{
template<typename TValue, typename TStackAllocator>
struct StlStackAllocator
{
public:
    using value_type = TValue;
private:
    TStackAllocator* base_;
public:
    explicit StlStackAllocator(TStackAllocator& base) MIJIN_NOEXCEPT : base_(&base) {}
    template<typename TOtherValue>
    StlStackAllocator(const StlStackAllocator<TOtherValue, TStackAllocator>& other) MIJIN_NOEXCEPT : base_(other.base) {}

    template<typename TOtherValue>
    StlStackAllocator& operator=(const StlStackAllocator<TOtherValue, TStackAllocator>& other) MIJIN_NOEXCEPT
    {
        base_ = other.base_;
        return *this;
    }

    auto operator<=>(const StlStackAllocator&) const noexcept = default;

    [[nodiscard]]
    TValue* allocate(std::size_t count) const
    {
        return static_cast<TValue*>(base_->allocate(alignof(TValue), count * sizeof(TValue)));
    }

    void deallocate(TValue* /* ptr */, std::size_t /* count */) const MIJIN_NOEXCEPT {}
};

template<std::size_t chunkSize = 4096, template<typename> typename TBacking = MIJIN_DEFAULT_ALLOCATOR> requires (allocator_tmpl<TBacking>)
class StackAllocator
{
public:
    using backing_t = TBacking<void>;
    static constexpr std::size_t ACTUAL_CHUNK_SIZE = chunkSize - sizeof(void*) - sizeof(std::size_t);

    template<typename T>
    using stl_allocator_t = StlStackAllocator<T, StackAllocator<chunkSize, TBacking>>;
private:
    struct Chunk
    {
        std::array<std::byte, ACTUAL_CHUNK_SIZE> data;
        Chunk* next;
        std::size_t allocated;
    };
    [[no_unique_address]] TBacking<Chunk> backing_;
    Chunk* firstChunk_ = nullptr;
public:
    StackAllocator() MIJIN_NOEXCEPT_IF(std::is_nothrow_default_constructible_v<backing_t>) = default;
    explicit StackAllocator(backing_t backing) MIJIN_NOEXCEPT_IF((std::is_nothrow_constructible_v<TBacking<Chunk>, backing_t&&>))
        : backing_(std::move(backing)) {}
    StackAllocator(const StackAllocator&) = delete;
    StackAllocator(StackAllocator&& other) MIJIN_NOEXCEPT_IF(std::is_nothrow_move_constructible_v<TBacking<Chunk>>)
        : firstChunk_(std::exchange(other.firstChunk_, nullptr)) {}
    ~StackAllocator() noexcept
    {
        Chunk* chunk = firstChunk_;
        while (chunk != nullptr)
        {
            Chunk* nextChunk = firstChunk_->next;
            backing_.deallocate(chunk, 1);
            chunk = nextChunk;
        }
    }
    StackAllocator& operator=(const StackAllocator&) = delete;
    StackAllocator& operator=(StackAllocator&& other) MIJIN_NOEXCEPT_IF(std::is_nothrow_move_assignable_v<TBacking<Chunk>>)
    {
        if (this != &other)
        {
            backing_ = std::move(other.backing_);
            firstChunk_ = std::exchange(other.firstChunk_, nullptr);
        }
        return *this;
    }

    void* allocate(std::size_t alignment, std::size_t size)
    {
        // first check if this can ever fit
        if (size > ACTUAL_CHUNK_SIZE)
        {
            return nullptr;
        }

        // then try to find space in the current chunks
        for (Chunk* chunk = firstChunk_; chunk != nullptr; chunk = chunk->next)
        {
            const std::size_t remaining = ACTUAL_CHUNK_SIZE - chunk->allocated;
            if (remaining < size)
            {
                continue;
            }

            std::byte* start = &chunk->data[chunk->allocated];
            std::byte* pos = mijin::alignUp(start, alignment);

            const std::ptrdiff_t alignmentBytes = pos - start;
            const std::size_t combinedSize = size + alignmentBytes;

            if (remaining < combinedSize)
            {
                continue;
            }

            chunk->allocated += combinedSize;
            return pos;
        }

        // no free space in any chunk? allocate a new one
        Chunk* newChunk = backing_.allocate(1);
        if (newChunk == nullptr)
        {
            return nullptr;
        }
        initAndAddChunk(newChunk);

        // now try with the new chunk
        std::byte* start = newChunk->data.data();
        std::byte* pos = mijin::alignUp(start, alignment);

        const std::ptrdiff_t alignmentBytes = pos - start;
        const std::size_t combinedSize = size + alignmentBytes;

        // doesn't fit (due to alignment), time to give up
        if (ACTUAL_CHUNK_SIZE < combinedSize)
        {
            return nullptr;
        }

        newChunk->allocated = combinedSize;
        return pos;
    }

    void reset() noexcept
    {
        for (Chunk* chunk = firstChunk_; chunk != nullptr; chunk = chunk->next)
        {
            chunk->allocated = 0;
        }
    }

    bool createChunks(std::size_t count)
    {
        if (count == 0)
        {
            return true;
        }
        Chunk* newChunks = backing_.allocate(count);
        if (newChunks == nullptr)
        {
            return false;
        }

        // reverse so the chunks are chained from 0 to count (new chunks are inserted in front)
        for (std::size_t pos = count; pos > 0; --pos)
        {
            initAndAddChunk(&newChunks[pos-1]);
        }
        return true;
    }

    template<typename T>
    stl_allocator_t<T> makeStlAllocator() MIJIN_NOEXCEPT
    {
        return stl_allocator_t<T>(*this);
    }
private:
    void initAndAddChunk(Chunk* newChunk) noexcept
    {
        ::new (newChunk) Chunk();

        // put it in the front
        newChunk->next = firstChunk_;
        firstChunk_ = newChunk;
    }
};
} // namespace mijin

#endif // !defined(MIJIN_MEMORY_STACK_ALLOCATOR_HPP_INCLUDED)