#include "XCEngine/RHI/D3D12/D3D12DescriptorHeap.h"

namespace XCEngine {
namespace RHI {

D3D12DescriptorHeap::D3D12DescriptorHeap()
    : m_type(DescriptorHeapType::CBV_SRV_UAV)
    , m_numDescriptors(0)
    , m_descriptorSize(0)
    , m_shaderVisible(false) {
}

D3D12DescriptorHeap::~D3D12DescriptorHeap() {
    Shutdown();
}

bool D3D12DescriptorHeap::Initialize(const DescriptorPoolDesc& desc) {
    return Initialize(
        static_cast<ID3D12Device*>(desc.device),
        desc.type,
        desc.descriptorCount,
        desc.shaderVisible
    );
}

bool D3D12DescriptorHeap::Initialize(ID3D12Device* device, DescriptorHeapType type, uint32_t numDescriptors, bool shaderVisible) {
    D3D12_DESCRIPTOR_HEAP_DESC desc = {};
    desc.Type = ToD3D12(type);
    desc.NumDescriptors = numDescriptors;
    desc.Flags = shaderVisible ? D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE : D3D12_DESCRIPTOR_HEAP_FLAG_NONE;
    desc.NodeMask = 0;

    HRESULT hResult = device->CreateDescriptorHeap(&desc, IID_PPV_ARGS(&m_descriptorHeap));
    if (FAILED(hResult)) {
        return false;
    }

    m_type = type;
    m_numDescriptors = numDescriptors;
    m_shaderVisible = shaderVisible;
    m_descriptorSize = device->GetDescriptorHandleIncrementSize(ToD3D12(type));

    return true;
}

void D3D12DescriptorHeap::Shutdown() {
    m_descriptorHeap.Reset();
}

D3D12_CPU_DESCRIPTOR_HANDLE D3D12DescriptorHeap::GetCPUDescriptorHandleForHeapStart() const {
    return m_descriptorHeap->GetCPUDescriptorHandleForHeapStart();
}

D3D12_GPU_DESCRIPTOR_HANDLE D3D12DescriptorHeap::GetGPUDescriptorHandleForHeapStart() const {
    return m_descriptorHeap->GetGPUDescriptorHandleForHeapStart();
}

CPUDescriptorHandle D3D12DescriptorHeap::GetCPUDescriptorHandle(uint32_t index) {
    D3D12_CPU_DESCRIPTOR_HANDLE handle = m_descriptorHeap->GetCPUDescriptorHandleForHeapStart();
    handle.ptr += index * m_descriptorSize;
    CPUDescriptorHandle result;
    result.ptr = handle.ptr;
    return result;
}

GPUDescriptorHandle D3D12DescriptorHeap::GetGPUDescriptorHandle(uint32_t index) {
    D3D12_GPU_DESCRIPTOR_HANDLE handle = m_descriptorHeap->GetGPUDescriptorHandleForHeapStart();
    handle.ptr += index * m_descriptorSize;
    GPUDescriptorHandle result;
    result.ptr = handle.ptr;
    return result;
}

DescriptorHeapType D3D12DescriptorHeap::GetType() const {
    return m_type;
}

uint32_t D3D12DescriptorHeap::GetDescriptorCount() const {
    return m_numDescriptors;
}

D3D12_DESCRIPTOR_HEAP_DESC D3D12DescriptorHeap::CreateDesc(DescriptorHeapType type, uint32_t numDescriptors, bool shaderVisible) {
    D3D12_DESCRIPTOR_HEAP_DESC desc = {};
    desc.Type = ToD3D12(type);
    desc.NumDescriptors = numDescriptors;
    desc.Flags = shaderVisible ? D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE : D3D12_DESCRIPTOR_HEAP_FLAG_NONE;
    desc.NodeMask = 0;
    return desc;
}

} // namespace RHI
} // namespace XCEngine