#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_PARAM_FIRST(const char* name = nullptr))
        {
            Assert(!Arena);
            static_assert(AllowRealloc == false);

            DataFirst = DataLast = nullptr;
            Count                = 0;

            JULIET_DEBUG_ONLY(Name = name ? name : Name;)
            ArenaParams params{ .AllowRealloc = AllowRealloc JULIET_DEBUG_PARAM(.CanReserveMore = false) };
            Arena         = ArenaAllocate(params JULIET_DEBUG_PARAM(Name));
            InternalArena = true;

            Reserve(ReserveSize);
        }

        void Create(NonNullPtr<Arena> arena JULIET_DEBUG_PARAM(const char* name = nullptr))
        {
            Assert(!Arena);

            JULIET_DEBUG_ONLY(Name = name ? name : Name;)

            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 JULIET_DEBUG_PARAM(Name));
                }
                else if constexpr (AllowRealloc)
                {
                    if (!InternalArena)
                    {
                        DataFirst = Data =
                            static_cast<Type*>(ArenaReallocate(Arena, Data, Capacity * sizeof(Type), newCapacity * sizeof(Type),
                                                               AlignOf(Type), true JULIET_DEBUG_PARAM("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* buffer, size_t amount)
        {
            Assert(Arena);

            if (Count + amount > Capacity)
            {
                if (Capacity == 0 && Count + amount < ReserveSize)
                {
                    Reserve(ReserveSize);
                }
                else
                {
                    size_t newCapacity = Max(Capacity * 2, AlignPow2(Capacity + amount, AlignOf(Type)));
                    Reserve(newCapacity);
                }
            }

            Type* dst = Data + Count;
            MemCopy(dst, buffer, amount * sizeof(Type));

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

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

        bool IsEmpty() const { return Count == 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* Front() { return DataFirst; }
        Type* Last() { return DataLast; }
        Type* Back() { 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;
        JULIET_DEBUG_ONLY(const char* Name = "VectorArena";)
    };
    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
