#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"

#if !defined(MIJIN_STACK_ALLOCATOR_DEBUG)
#if defined(MIJIN_DEBUG)
#define MIJIN_STACK_ALLOCATOR_DEBUG 1
#else
#define MIJIN_STACK_ALLOCATOR_DEBUG 0
#endif
#endif // !defined(MIJIN_STACK_ALLOCATOR_DEBUG)

#if MIJIN_STACK_ALLOCATOR_DEBUG > 1
#include <print>
#include <unordered_map>
#include "../debug/stacktrace.hpp"
#endif

namespace mijin
{
template<typename TValue, typename TStackAllocator>
class 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
    {
        void* result = base_->allocate(alignof(TValue), count * sizeof(TValue));
#if MIJIN_STACK_ALLOCATOR_DEBUG == 1
        ++base_->numAllocations_;
#elif MIJIN_STACK_ALLOCATOR_DEBUG > 1
        base_->activeAllocations_.emplace(result, captureStacktrace(1));
#endif
        return static_cast<TValue*>(result);
    }

    void deallocate([[maybe_unused]] TValue* ptr, std::size_t /* count */) const MIJIN_NOEXCEPT
    {
#if MIJIN_STACK_ALLOCATOR_DEBUG == 1
        MIJIN_ASSERT(base_->numAllocations_ > 0, "Unbalanced allocations in stack allocators!");
        --base_->numAllocations_;
#elif MIJIN_STACK_ALLOCATOR_DEBUG > 1
        auto it = base_->activeAllocations_.find(ptr);
        if (it != base_->activeAllocations_.end())
        {
            base_->activeAllocations_.erase(it);
        }
        else
        {
            MIJIN_ERROR("Deallocating invalid pointer from StackAllocator.");
        }
#endif
    }

    template<typename TOtherValue, typename TOtherAllocator>
    friend class StlStackAllocator;
};

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;
#if MIJIN_STACK_ALLOCATOR_DEBUG == 1
    std::size_t numAllocations_ = 0;
#elif MIJIN_STACK_ALLOCATOR_DEBUG > 1
    // just for debugging, so we don't care what memory this uses...
    std::unordered_map<void*, Result<Stacktrace>> activeAllocations_;
#endif
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
    {
#if MIJIN_STACK_ALLOCATOR_DEBUG == 1
        MIJIN_ASSERT(numAllocations_ == 0, "Missing deallocation in StackAllocator!");
#elif MIJIN_STACK_ALLOCATOR_DEBUG > 1
        if (!activeAllocations_.empty())
        {
            std::println(stderr, "{} active allocations in StackAllocator when resetting!", activeAllocations_.size());
            for (const auto& [ptr, stack] : activeAllocations_)
            {
                if (stack.isError())
                {
                    std::println(stderr, "at {}, no stacktrace ({})", ptr, stack.getError().message);
                }
                else
                {
                    std::println(stderr, "at 0x{}:\n{}", ptr, stack.getValue());
                }
            }
            MIJIN_TRAP();
        }
#endif
        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;
    }

    template<typename TValue, typename TStackAllocator>
    friend class StlStackAllocator;
};
} // namespace mijin

#endif // !defined(MIJIN_MEMORY_STACK_ALLOCATOR_HPP_INCLUDED)