
#include <Core/Memory/Allocator.h>
#include <Core/Memory/Utils.h>
#include <Graphics/D3D12/D3D12Common.h>
#include <Graphics/D3D12/D3D12GraphicsDevice.h>
#include <Graphics/D3D12/D3D12Includes.h>

// TODO: Convert the whole file to memory arenas
namespace Juliet::D3D12::Internal
{
    namespace
    {
        constexpr size_t kStagingHeapDescriptorExpectedCount = 1024;

        void InitStagingDescriptorPool(NonNullPtr<D3D12DescriptorHeap> heap, NonNullPtr<D3D12StagingDescriptorPool> pool)
        {
            for (uint32 idx = 0; idx < kStagingHeapDescriptorExpectedCount; ++idx)
            {
                pool->FreeDescriptors[idx].Pool           = pool;
                pool->FreeDescriptors[idx].Heap           = heap.Get();
                pool->FreeDescriptors[idx].CpuHandleIndex = idx;
                pool->FreeDescriptors[idx].CpuHandle.ptr = heap->DescriptorHeapCPUStart.ptr + (idx * heap->DescriptorSize);
            }
        }

        bool ExtendStagingDescriptorPool(NonNullPtr<D3D12Driver> driver, D3D12StagingDescriptorPool& pool)
        {
            D3D12DescriptorHeap* heap = Internal::CreateDescriptorHeap(driver, pool.Arena, pool.Heaps[0]->HeapType,
                                                                       kStagingHeapDescriptorExpectedCount, true);
            if (!heap)
            {
                return false;
            }

            pool.HeapCount += 1;
            pool.Heaps = static_cast<D3D12DescriptorHeap**>(Realloc(pool.Heaps, pool.HeapCount * sizeof(D3D12DescriptorHeap*)));
            pool.Heaps[pool.HeapCount - 1] = heap;

            pool.FreeDescriptorCapacity += kStagingHeapDescriptorExpectedCount;
            pool.FreeDescriptorCount += kStagingHeapDescriptorExpectedCount;
            pool.FreeDescriptors = static_cast<D3D12StagingDescriptor*>(
                Realloc(pool.FreeDescriptors, pool.FreeDescriptorCapacity * sizeof(D3D12StagingDescriptor)));

            InitStagingDescriptorPool(heap, &pool);

            return true;
        }
    } // namespace

    D3D12StagingDescriptorPool* CreateStagingDescriptorPool(NonNullPtr<D3D12Driver> driver, D3D12_DESCRIPTOR_HEAP_TYPE type)
    {
        Arena* arena = ArenaAllocate({} JULIET_DEBUG_ONLY(, "Staging Descriptors"));

        D3D12DescriptorHeap* heap = CreateDescriptorHeap(driver, arena, type, kStagingHeapDescriptorExpectedCount, true);
        if (!heap)
        {
            ArenaRelease(arena);
            return nullptr;
        }

        auto pool   = static_cast<D3D12StagingDescriptorPool*>(Calloc(1, sizeof(D3D12StagingDescriptorPool)));
        pool->Arena = arena;

        // First create the heaps
        pool->HeapCount = 1;
        pool->Heaps     = static_cast<D3D12DescriptorHeap**>(Malloc(sizeof(D3D12DescriptorHeap*)));
        pool->Heaps[0]  = heap;

        pool->FreeDescriptorCapacity = kStagingHeapDescriptorExpectedCount;
        pool->FreeDescriptorCount    = kStagingHeapDescriptorExpectedCount;
        pool->FreeDescriptors =
            static_cast<D3D12StagingDescriptor*>(Malloc(kStagingHeapDescriptorExpectedCount * sizeof(D3D12StagingDescriptor)));

        InitStagingDescriptorPool(heap, pool);

        return pool;
    }

    bool AssignStagingDescriptor(NonNullPtr<D3D12Driver> driver, D3D12_DESCRIPTOR_HEAP_TYPE type, D3D12StagingDescriptor& outDescriptor)
    {
        // TODO: Make it thread safe
        D3D12StagingDescriptor*     descriptor = nullptr;
        D3D12StagingDescriptorPool* pool       = driver->StagingDescriptorPools[type];

        if (pool->FreeDescriptorCount == 0)
        {
            if (!ExtendStagingDescriptorPool(driver, *pool))
            {
                return false;
            }
        }

        descriptor = &pool->FreeDescriptors[pool->FreeDescriptorCount - 1];
        MemCopy(&outDescriptor, descriptor, sizeof(D3D12StagingDescriptor));
        pool->FreeDescriptorCount -= 1;

        return true;
    }

    void ReleaseStagingDescriptor(NonNullPtr<D3D12Driver> /*driver*/, D3D12StagingDescriptor& cpuDescriptor)
    {
        D3D12StagingDescriptorPool* pool = cpuDescriptor.Pool;

        if (pool != nullptr)
        {
            MemCopy(&pool->FreeDescriptors[pool->FreeDescriptorCount], &cpuDescriptor, sizeof(D3D12StagingDescriptor));
            pool->FreeDescriptorCount += 1;
        }
    }

    void DestroyStagingDescriptorPool(NonNullPtr<D3D12StagingDescriptorPool> pool)
    {
        for (uint32 i = 0; i < pool->HeapCount; i += 1)
        {
            DestroyDescriptorHeap(pool->Heaps[i]);
        }

        Free(pool->Heaps);
        Free(pool->FreeDescriptors);

        ArenaRelease(pool->Arena);
        Free(pool.Get());
    }
} // namespace Juliet::D3D12::Internal
