#include <Core/Memory/Allocator.h>
#include <Core/Memory/EngineArena.h>
#include <Graphics/D3D12/D3D12DescriptorHeap.h>
#include <Graphics/D3D12/D3D12GraphicsDevice.h>
#include <Graphics/D3D12/D3D12Utils.h>

namespace Juliet::D3D12::Internal
{
    D3D12DescriptorHeap* CreateDescriptorHeap(NonNullPtr<D3D12Driver> driver, D3D12_DESCRIPTOR_HEAP_TYPE type, uint32 count, bool isStaging)
    {
        ID3D12DescriptorHeap* handle;

        auto heap = ArenaPushType<D3D12DescriptorHeap>(GetEngineArena(), ConstString("D3D12DescriptorHeap"));
        if (!heap)
        {
            return nullptr;
        }

        heap->CurrentDescriptorIndex = 0;
        heap->FreeIndicesCapacity    = 16;
        heap->FreeIndicesCount       = 0;
        heap->FreeIndices            = ArenaPushArray<uint32>(GetEngineArena(), heap->FreeIndicesCapacity,
                                                              ConstString("D3D12DescriptorHeap/FreeIndices"));

        D3D12_DESCRIPTOR_HEAP_DESC heapDesc;
        heapDesc.NumDescriptors = count;
        heapDesc.Type           = type;
        heapDesc.Flags    = isStaging ? D3D12_DESCRIPTOR_HEAP_FLAG_NONE : D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE;
        heapDesc.NodeMask = 0;

        HRESULT result = driver->D3D12Device->CreateDescriptorHeap(&heapDesc, IID_ID3D12DescriptorHeap, (void**)&handle);
        if (FAILED(result))
        {
            LogError(driver->D3D12Device, "Failed to create descriptor heap!", result);
            DestroyDescriptorHeap(heap);
            return nullptr;
        }

        heap->Handle                 = handle;
        heap->HeapType               = type;
        heap->MaxDescriptors         = count;
        heap->Staging                = isStaging;
        heap->DescriptorSize         = driver->D3D12Device->GetDescriptorHandleIncrementSize(type);
        heap->DescriptorHeapCPUStart = handle->GetCPUDescriptorHandleForHeapStart();
        if (!isStaging)
        {
            heap->DescriptorHeapGPUStart = handle->GetGPUDescriptorHandleForHeapStart();
        }

        return heap;
    }

    void DestroyDescriptorHeap(NonNullPtr<D3D12DescriptorHeap> heap)
    {
        if (heap->Handle)
        {
            heap->Handle->Release();
        }
        SafeFree(heap->FreeIndices);
        Free(heap.Get());
    }

    bool AssignDescriptor(D3D12DescriptorHeap* heap, D3D12Descriptor& outDescriptor)
    {
        uint32 index = UINT32_MAX;

        if (heap->FreeIndicesCount > 0)
        {
            heap->FreeIndicesCount -= 1;
            index = heap->FreeIndices[heap->FreeIndicesCount];
        }
        else if (heap->CurrentDescriptorIndex < heap->MaxDescriptors)
        {
            index = heap->CurrentDescriptorIndex;
            heap->CurrentDescriptorIndex++;
        }
        else
        {
            Assert(false, "Descriptor Heap Full!");
            return false;
        }

        outDescriptor.Heap      = heap;
        outDescriptor.Index     = index;
        outDescriptor.CpuHandle = heap->DescriptorHeapCPUStart;
        outDescriptor.CpuHandle.ptr += heap->DescriptorSize * index;
        outDescriptor.GpuHandle = heap->DescriptorHeapGPUStart;
        outDescriptor.GpuHandle.ptr += heap->DescriptorSize * index;

        return true;
    }

    void ReleaseDescriptor(const D3D12Descriptor& descriptor)
    {
        if (descriptor.Index == UINT32_MAX || descriptor.Heap == nullptr)
        {
            return;
        }

        D3D12DescriptorHeap* heap = descriptor.Heap;

        if (heap->FreeIndicesCount >= heap->FreeIndicesCapacity)
        {
            heap->FreeIndicesCapacity *= 2;
            heap->FreeIndices = static_cast<uint32*>(Realloc(heap->FreeIndices, heap->FreeIndicesCapacity * sizeof(uint32)));
        }

        heap->FreeIndices[heap->FreeIndicesCount] = descriptor.Index;
        heap->FreeIndicesCount++;
    }

    D3D12DescriptorHeap* AcquireSamplerHeapFromPool(NonNullPtr<D3D12Driver> d3d12Driver, DescriptorHeapCreator creator)
    {
        D3D12DescriptorHeapPool* pool  = &d3d12Driver->SamplerHeapPool;
        uint32                   count = GPUDriver::kSampler_HeapDescriptorCount;

        if (pool->Count > 0)
        {
            pool->Count -= 1;
            return pool->Heaps[pool->Count];
        }

        // Pool is exhausted (all heaps are in use). We need to create a new one.
        // The pool array might be full of "Used" heaps, or it might have empty slots (nullptr) if we expanded before.
        // We look for a free slot to store the new heap.
        size_t freeSlotIndex = SIZE_MAX;

        // Scan for nullptr in the entire array.
        // Optimization: active heaps are likely at the beginning, but since Count=0, "Used" heaps are 0..Capacity?
        // Actually, if Count=0, ALL slots 0..Capacity-1 are effectively "Used" (unless some are nullptr).
        for (size_t i = 0; i < pool->Capacity; ++i)
        {
            if (pool->Heaps[i] == nullptr)
            {
                freeSlotIndex = i;
                break;
            }
        }

        if (freeSlotIndex == SIZE_MAX)
        {
            // No free slot, expand
            size_t oldCapacity = pool->Capacity;
            pool->Capacity     = pool->Capacity == 0 ? 1 : pool->Capacity * 2;

            pool->Heaps = ArenaRealloc<D3D12DescriptorHeap*>(GetEngineArena(), pool->Heaps, oldCapacity, pool->Capacity,
                                                             ConstString("DescriptorHeapArray"));

            // Initialize new slots to nullptr
            for (size_t i = oldCapacity; i < pool->Capacity; ++i)
            {
                pool->Heaps[i] = nullptr;
            }
            freeSlotIndex = oldCapacity;
        }

        // Create the new heap at the free slot
        pool->Heaps[freeSlotIndex] = creator(d3d12Driver, D3D12_DESCRIPTOR_HEAP_TYPE_SAMPLER, count, false);

        // Now we have a new heap at freeSlotIndex. Ideally we want to return it.
        // But to maintain the "Stack" invariant (Available are at 0..Count-1),
        // and "Used" are at "Count..Capacity-1".
        // Currently Count=0. So "Used" is 0..Capacity-1.
        // freeSlotIndex is inside "Used" range.
        // We can just return it. It is already in the "Used" partition.
        // However, we want to ensure it doesn't get lost or overwritten.
        // Since we didn't increment Count, it stays in "Used".

        return pool->Heaps[freeSlotIndex];
    }

    void ReturnSamplerHeapToPool(NonNullPtr<D3D12Driver> d3d12Driver, D3D12DescriptorHeap* heap)
    {
        if (heap == nullptr)
        {
            return;
        }

        Assert(heap->HeapType == D3D12_DESCRIPTOR_HEAP_TYPE_SAMPLER);
        D3D12DescriptorHeapPool* pool = &d3d12Driver->SamplerHeapPool;

        heap->CurrentDescriptorIndex = 0;

        // The heap is currently in the "Used" partition (index >= Count).
        // We want to move it to the "Available" partition (index < Count).
        // Strategy: Swap heap with the element at pool->Heaps[pool->Count], then increment Count.

        // First, find the heap index.
        size_t heapIndex = SIZE_MAX;
        for (size_t i = pool->Count; i < pool->Capacity; ++i)
        {
            if (pool->Heaps[i] == heap)
            {
                heapIndex = i;
                break;
            }
        }

        if (heapIndex != SIZE_MAX)
        {
            D3D12DescriptorHeap* temp = pool->Heaps[pool->Count];
            pool->Heaps[pool->Count]  = pool->Heaps[heapIndex];
            pool->Heaps[heapIndex]    = temp;
            pool->Count += 1;
        }
        else
        {
            // Should not happen if the logic is correct.
            // Maybe it was not in the pool? (e.g. error scenario).
            // In that case, we should probably add it?
            // But per constraints, we assume it was allocated via pool.
            Assert(false, "Returning a heap that was not found in the pool!");
        }
    }
} // namespace Juliet::D3D12::Internal
