#pragma once

#if !defined(MIJIN_MEMORY_STACK_ALLOCATOR_HPP_INCLUDED)
#define MIJIN_MEMORY_STACK_ALLOCATOR_HPP_INCLUDED 1

#include <memory>
#include <utility>
#include "../debug/assert.hpp"
#include "../internal/common.hpp"
#include "../util/align.hpp"
#include "../util/concepts.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;
};

namespace impl
{
struct StackAllocatorSnapshotData
{
    struct ChunkSnapshot
    {
        std::size_t allocated;
    };
    std::size_t numChunks;
#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
    ChunkSnapshot chunks[1]; // may be bigger than that
};
}

class StackAllocatorSnapshot
{
private:
    impl::StackAllocatorSnapshotData* data = nullptr;

    impl::StackAllocatorSnapshotData* operator->() const MIJIN_NOEXCEPT { return data; }


    template<std::size_t chunkSize, template<typename> typename TBacking> requires (allocator_tmpl<TBacking>)
    friend class StackAllocator;
};

template<typename TStackAllocator>
class StackAllocatorScope
{
private:
    TStackAllocator* allocator_;
    StackAllocatorSnapshot snapshot_;
public:
    explicit StackAllocatorScope(TStackAllocator& allocator) : allocator_(&allocator), snapshot_(allocator_->createSnapshot()) {}
    StackAllocatorScope(const StackAllocatorScope&) = delete;
    StackAllocatorScope(StackAllocatorScope&&) = delete;
    ~StackAllocatorScope() MIJIN_NOEXCEPT
    {
        allocator_->restoreSnapshot(snapshot_);
    }

    StackAllocatorScope& operator=(const StackAllocatorScope&) = delete;
    StackAllocatorScope& operator=(StackAllocatorScope&&) = delete;
};

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);
    }

    [[nodiscard]]
    std::size_t getNumChunks() const MIJIN_NOEXCEPT
    {
        std::size_t num = 0;
        for (Chunk* chunk = firstChunk_; chunk != nullptr; chunk = chunk->next)
        {
            ++num;
        }
        return num;
    }

    [[nodiscard]]
    StackAllocatorSnapshot createSnapshot()
    {
        if (firstChunk_ == nullptr)
        {
            return {};
        }

        using impl::StackAllocatorSnapshotData;

        std::size_t numChunks = getNumChunks();
        std::size_t snapshotSize = calcSnapshotSize(numChunks);
        Chunk* prevFirst = firstChunk_;
        StackAllocatorSnapshotData* snapshotData = static_cast<StackAllocatorSnapshotData*>(allocate(alignof(StackAllocatorSnapshotData), snapshotSize));
        if (snapshotData == nullptr)
        {
            // couldn't allocate the snapshot
            return {};
        }
        ::new (snapshotData) StackAllocatorSnapshotData;
        StackAllocatorSnapshot snapshot;
        snapshot.data = snapshotData;
        if (firstChunk_ != prevFirst)
        {
            // new chunk has been added, adjust the snapshot
            // the snapshot must be inside the new chunk (but not necessarily at the very beginning, due to alignment)
            MIJIN_ASSERT(static_cast<const void*>(snapshot.data) == alignUp(firstChunk_->data.data(), alignof(StackAllocatorSnapshot)), "Snapshot not where it was expected.");

            // a chunk might be too small for the snapshot if we grow it (unlikely, but not impossible)
            if (ACTUAL_CHUNK_SIZE - firstChunk_->allocated < MIJIN_STRIDEOF(StackAllocatorSnapshotData::ChunkSnapshot)) [[unlikely]]
            {
                firstChunk_->allocated = 0;
                return {};
            }

            // looking good, adjust the numbers
            ++numChunks;
            snapshotSize += MIJIN_STRIDEOF(StackAllocatorSnapshotData::ChunkSnapshot);
            firstChunk_->allocated += MIJIN_STRIDEOF(StackAllocatorSnapshotData::ChunkSnapshot);
        }
        // now fill out the struct
        snapshot->numChunks = numChunks;
#if MIJIN_STACK_ALLOCATOR_DEBUG == 1
        snapshot->numAllocations_ = numAllocations_;
#elif MIJIN_STACK_ALLOCATOR_DEBUG > 1
        // just for debugging, so we don't care what memory this uses...
        snapshot->activeAllocations_ = activeAllocations_;
#endif

        std::size_t pos = 0;
        for (Chunk* chunk = firstChunk_; chunk != nullptr; chunk = chunk->next, ++pos)
        {
            snapshot->chunks[pos].allocated = chunk->allocated;
        }
        return snapshot;
    }

    void restoreSnapshot(StackAllocatorSnapshot snapshot) MIJIN_NOEXCEPT
    {
        if (snapshot.data == nullptr)
        {
            return;
        }

        const std::size_t numChunks = getNumChunks();
        MIJIN_ASSERT_FATAL(snapshot->numChunks <= numChunks, "Snapshot contains more chunks than the allocator!");

#if MIJIN_STACK_ALLOCATOR_DEBUG == 1
        MIJIN_ASSERT(snapshot->numAllocations_ >= numAllocations_, "Missing deallocation in StackAllocator!");
#elif MIJIN_STACK_ALLOCATOR_DEBUG > 1
        // TODO: compare and print changes
        unsigned numMismatches = 0;
        for (const auto& [ptr, stack] : activeAllocations_)
        {
            if (snapshot->activeAllocations_.contains(ptr))
            {
                continue;
            }
            ++numMismatches;

            if (stack.isError())
            {
                std::println(stderr, "Missing deallocation at {}, no stacktrace ({})", ptr, stack.getError().message);
            }
            else
            {
                std::println(stderr, "Missing deallocation at 0x{}:\n{}", ptr, stack.getValue());
            }
        }
#if 0 // deallocating more than expected shouldn't be a problem
        for (const auto& [ptr, stack] : snapshot->activeAllocations_)
        {
            if (activeAllocations_.contains(ptr))
            {
                continue;
            }
            ++numMismatches;

            if (stack.isError())
            {
                std::println(stderr, "Unexpected deallocation at {}, no stacktrace ({})", ptr, stack.getError().message);
            }
            else
            {
                std::println(stderr, "Unexpected deallocation at 0x{}:\n{}", ptr, stack.getValue());
            }
        }
#endif
        if (numMismatches > 0)
        {
            std::println(stderr, "{} mismatched deallocations when restoring stack allocator snapshot.", numMismatches);
            MIJIN_TRAP();
        }
#endif
        // if we allocated new chunks since the snapshot, these are completely empty now
        const std::size_t emptyChunks = numChunks - snapshot->numChunks;

        Chunk* chunk = firstChunk_;
        for (std::size_t idx = 0; idx < emptyChunks; ++idx)
        {
            chunk->allocated = 0;
            chunk = chunk->next;
        }

        // the other values are in the snapshot
        for (std::size_t idx = 0; idx < snapshot->numChunks; ++idx)
        {
            chunk->allocated = snapshot->chunks[idx].allocated;
            chunk = chunk->next;
        }
        MIJIN_ASSERT(chunk == nullptr, "Something didn't add up.");

        // finally free the space for the snapshot itself
        Chunk* snapshotChunk = findChunk(snapshot.data);
        MIJIN_ASSERT_FATAL(snapshotChunk != nullptr, "Snapshot not in chunks?");
        snapshotChunk->allocated -= calcSnapshotSize(snapshot->numChunks); // note: this might miss the alignment bytes of the snapshot, but that should be fine
        snapshot.data->~StackAllocatorSnapshotData();
    }
private:
    void initAndAddChunk(Chunk* newChunk) noexcept
    {
        ::new (newChunk) Chunk();

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

    bool isInChunk(const void* address, const Chunk& chunk) const MIJIN_NOEXCEPT
    {
        const std::byte* asByte = static_cast<const std::byte*>(address);
        return asByte >= chunk.data.data() && asByte < chunk.data.data() + ACTUAL_CHUNK_SIZE;
    }

    Chunk* findChunk(const void* address) const MIJIN_NOEXCEPT
    {
        for (Chunk* chunk = firstChunk_; chunk != nullptr; chunk = chunk->next)
        {
            if (isInChunk(address, *chunk))
            {
                return chunk;
            }
        }
        return nullptr;
    }

    static std::size_t calcSnapshotSize(std::size_t numChunks) MIJIN_NOEXCEPT
    {
        return sizeof(impl::StackAllocatorSnapshotData) + ((numChunks - 1) * MIJIN_STRIDEOF(impl::StackAllocatorSnapshotData::ChunkSnapshot));
    }

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

#endif // !defined(MIJIN_MEMORY_STACK_ALLOCATOR_HPP_INCLUDED)