#pragma once

#include <Core/Common/NonNullPtr.h>
#include <Core/Memory/MemoryArena.h>
#include <vector>

namespace Juliet
{
    template <typename Type, size_t ReserveSize = 16, bool AllowRealloc = false>
    struct VectorArena
    {
        void Create(JULIET_DEBUG_ONLY(const char* name = "VectorArena"))
        {
            Assert(!Arena);
            static_assert(AllowRealloc == false);

            DataFirst = DataLast = nullptr;
            Count                = 0;
            ArenaParams params{ .AllowRealloc = AllowRealloc JULIET_DEBUG_ONLY(, .CanReserveMore = false) };
            Arena         = ArenaAllocate(params JULIET_DEBUG_ONLY(, name));
            InternalArena = true;

            Reserve(ReserveSize);
        }

        void Create(NonNullPtr<Arena> arena)
        {
            Assert(!Arena);

            DataFirst = DataLast = nullptr;
            Count                = 0;
            Arena                = arena.Get();
            InternalArena        = false;

            Reserve(ReserveSize);
        }

        void Destroy()
        {
            if (InternalArena)
            {
                ArenaRelease(Arena);
            }
            DataFirst = DataLast = nullptr;
            Count                = 0;
            Capacity             = 0;
            Arena                = nullptr;
        }

        void Reserve(size_t newCapacity)
        {
            if (newCapacity > Capacity)
            {
                if (Data == nullptr)
                {
                    Data = ArenaPushArray<Type>(Arena, newCapacity);
                }
                else if (AllowRealloc && !InternalArena)
                {
                    DataFirst = Data = static_cast<Type*>(ArenaReallocate(Arena, Data, Capacity * sizeof(Type),
                                                                          newCapacity * sizeof(Type), AlignOf(Type), true
                                                                          JULIET_DEBUG_ONLY(, "VectorRealloc")));
                    DataLast = Data + Count - 1;
                }
                Capacity = newCapacity;
            }
        }

        void Resize(size_t newCount)
        {
            if (newCount == Count)
            {
                return;
            }

            if (newCount > Capacity)
            {
                Reserve(newCount);
            }

            Count = newCount;

            if (Count > 0)
            {
                DataFirst = Data;
                DataLast  = Data + Count - 1;
            }
            else
            {
                DataFirst = DataLast = nullptr;
            }
        }

        void PushBack(const Type& value)
        {
            Assert(Arena);

            if (Count + 1 > Capacity)
            {
                Reserve(Capacity == 0 ? ReserveSize : Capacity * 2);
            }

            Type* entry = Data + Count;
            *entry      = value;

            if (Count == 0)
            {
                DataFirst = entry;
            }
            DataLast = entry;
            ++Count;
        }

        void PushBack(Type&& value)
        {
            Assert(Arena);

            if (Count + 1 > Capacity)
            {
                Reserve(Capacity == 0 ? ReserveSize : Capacity * 2);
            }

            Type* entry = Data + Count;
            *entry      = std::move(value);

            if (Count == 0)
            {
                DataFirst = entry;
            }
            DataLast = entry;
            ++Count;
        }

        void RemoveAtFast(index_t index)
        {
            Assert(Arena);
            Assert(index < Count);
            Assert(Count > 0);

            Type* elementAdr = DataFirst + index;

            // Swap DataLast and element
            if (DataLast != elementAdr)
            {
                Swap(DataLast, elementAdr);
            }

            --DataLast;
            --Count;

            if (Count == 0)
            {
                DataFirst = DataLast = nullptr;
            }
        }

        void Clear()
        {
            Assert(Arena);

            ArenaClear(Arena);
            DataFirst = DataLast = nullptr;
            Data                 = nullptr;
            Count                = 0;
            Capacity             = 0;
        }

        // C++ Accessors for loop supports and Index based access
        Type&       operator[](size_t index) { return DataFirst[index]; }
        const Type& operator[](size_t index) const { return DataFirst[index]; }

        Type* begin() { return DataFirst; }
        Type* end() { return DataFirst + Count; }

        const Type* begin() const { return DataFirst; }
        const Type* end() const { return DataFirst + Count; }

        Type* First() { return DataFirst; }
        Type* Last() { return DataLast; }

        size_t Size() const { return Count; }

        Arena* Arena;
        Type*  DataFirst;
        Type*  DataLast;
        Type*  Data;
        size_t Count;
        size_t Capacity;
        bool   InternalArena : 1;
    };
    static_assert(std::is_standard_layout_v<VectorArena<int>>,
                  "VectorArena must have a standard layout to remain POD-like.");
    static_assert(std::is_trivially_copyable_v<VectorArena<int>>,
                  "VectorArena must be trivially copyable (no custom destructors/assignment).");
} // namespace Juliet
