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

#if JULIET_DEBUG
    
namespace Juliet::D3D12::UnitTest
{
    using namespace Juliet::D3D12;
    using namespace Juliet::D3D12::Internal;

    static D3D12DescriptorHeap* MockCreateDescriptorHeap(NonNullPtr<D3D12Driver>, D3D12_DESCRIPTOR_HEAP_TYPE type, uint32, bool)
    {
        D3D12DescriptorHeap* heap = static_cast<D3D12DescriptorHeap*>(Calloc(1, sizeof(D3D12DescriptorHeap)));
        if (heap)
        {
            heap->HeapType = type;
            heap->MaxDescriptors = 128; 
            // Mock other fields as needed to pass verification
        }
        return heap;
    }

    void TestDescriptorHeapPool()
    {
        D3D12Driver driver = {};
        ZeroStruct(driver);
        // Ensure pool is empty
        driver.SamplerHeapPool.Count    = 0;
        driver.SamplerHeapPool.Capacity = 0;
        driver.SamplerHeapPool.Heaps    = nullptr;

        NonNullPtr<D3D12Driver> driverPtr(&driver);

        // 1. Acquire from empty -> Should create new
        D3D12DescriptorHeap* heap1 = AcquireSamplerHeapFromPool(driverPtr, MockCreateDescriptorHeap);
        Assert(heap1 != nullptr);
        Assert(driver.SamplerHeapPool.Count == 0);          // Used heaps are above Count
        Assert(driver.SamplerHeapPool.Capacity >= 1);
        Assert(driver.SamplerHeapPool.Heaps[0] == heap1);

        // 2. Return -> Should become available
        ReturnSamplerHeapToPool(driverPtr, heap1);
        Assert(driver.SamplerHeapPool.Count == 1);          // One available
        Assert(driver.SamplerHeapPool.Heaps[0] == heap1);   // Available is at index 0

        // 3. Acquire again -> Should reuse
        D3D12DescriptorHeap* heap2 = AcquireSamplerHeapFromPool(driverPtr, MockCreateDescriptorHeap);
        Assert(heap2 == heap1);                             // Reused
        Assert(driver.SamplerHeapPool.Count == 0);          // No available left

        // 4. Acquire another -> Should Expand and Create
        D3D12DescriptorHeap* heap3 = AcquireSamplerHeapFromPool(driverPtr, MockCreateDescriptorHeap);
        Assert(heap3 != nullptr);
        Assert(heap3 != heap1);
        Assert(driver.SamplerHeapPool.Count == 0);
        Assert(driver.SamplerHeapPool.Capacity >= 2);

        // Cleanup
        (void)heap1; (void)heap2; (void)heap3;
        Free(heap1);
        Free(heap3);

        // 5. Test Multiple Allocations (8 heaps)
        // Reset pool for clean test or continue? Let's reset to be sure of state, 
        // effectively leaking previous pool array but that's fine for mock test.
        // Actually, let's just continue using the pool, it has Capacity >= 2.
        
        // Return everything first to have clean slate if possible, but we already freed heap1 and heap3 mentally/physically.
        // The pool thinks heap3 is "Used" (outside Count). 
        // We freed it physically, but pool has a dangling pointer in Heaps[1].
        // This is dangerous if we reuse it.
        // So we MUST correct the pool state or reset it.
        // Let's reset the pool state for the 8-loop test.
        driver.SamplerHeapPool.Count = 0;
        driver.SamplerHeapPool.Capacity = 0;
        driver.SamplerHeapPool.Heaps = nullptr; // Leaks previous array

        const int kNumHeaps = 32;
        D3D12DescriptorHeap* heaps[kNumHeaps];

        // 5.1 Acquire 8
        for (int i = 0; i < kNumHeaps; ++i)
        {
            heaps[i] = AcquireSamplerHeapFromPool(driverPtr, MockCreateDescriptorHeap);
            Assert(heaps[i] != nullptr);
            // Verify uniqueness
            for (int j = 0; j < i; ++j)
            {
                Assert(heaps[i] != heaps[j]);
            }
        }
        Assert(driver.SamplerHeapPool.Capacity >= kNumHeaps);
        Assert(driver.SamplerHeapPool.Count == 0);

        // 5.2 Release 8
        for (int i = 0; i < kNumHeaps; ++i)
        {
            ReturnSamplerHeapToPool(driverPtr, heaps[i]);
        }
        Assert(driver.SamplerHeapPool.Count == kNumHeaps);

        // 5.3 Acquire 8 Again (Should Reuse)
        D3D12DescriptorHeap* heapsAgain[kNumHeaps];
        for (int i = 0; i < kNumHeaps; ++i)
        {
            heapsAgain[i] = AcquireSamplerHeapFromPool(driverPtr, MockCreateDescriptorHeap);
            
            // Verify it matches one of the original heaps (order might reverse due to stack-like behavior)
            bool found = false;
            for (int j = 0; j < kNumHeaps; ++j)
            {
                if (heapsAgain[i] == heaps[j]) 
                {
                    found = true; 
                    break;
                }
            }
            Assert(found);
             // Verify uniqueness in the new set
            for (int j = 0; j < i; ++j)
            {
                Assert(heapsAgain[i] != heapsAgain[j]);
            }
        }
        Assert(driver.SamplerHeapPool.Count == 0);
        // Capacity should not have increased if we reused
        Assert(driver.SamplerHeapPool.Capacity == kNumHeaps); // Or >= 8

        // Cleanup
        for (int i = 0; i < kNumHeaps; ++i)
        {
            Free(heaps[i]);
        }
        
        // Note: heap2 is heap1.
        // Pool array in driver leaked for test scope, acceptable.

        printf("DescriptorHeapPool tests passed.\n");
    }

}
#endif
