#include "XCEngine/RHI/D3D12/D3D12DescriptorSet.h"
#include "XCEngine/RHI/D3D12/D3D12DescriptorHeap.h"
#include "XCEngine/RHI/D3D12/D3D12ResourceView.h"
#include "XCEngine/RHI/D3D12/D3D12Sampler.h"

namespace XCEngine {
namespace RHI {

D3D12DescriptorSet::D3D12DescriptorSet()
    : m_heap(nullptr)
    , m_offset(0)
    , m_count(0)
    , m_bindingCount(0)
    , m_bindings(nullptr) {
}

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

bool D3D12DescriptorSet::Initialize(D3D12DescriptorHeap* heap, uint32_t offset, uint32_t count, const DescriptorSetLayoutDesc& layout) {
    m_heap = heap;
    m_offset = offset;
    m_count = count;
    m_bindingCount = layout.bindingCount;
    
    if (layout.bindingCount > 0 && layout.bindings != nullptr) {
        m_bindings = new DescriptorSetLayoutBinding[layout.bindingCount];
        for (uint32_t i = 0; i < layout.bindingCount; ++i) {
            m_bindings[i] = layout.bindings[i];
        }
    }
    
    return true;
}

void D3D12DescriptorSet::Shutdown() {
    m_heap = nullptr;
    m_offset = 0;
    m_count = 0;
    m_bindingCount = 0;
    if (m_bindings != nullptr) {
        delete[] m_bindings;
        m_bindings = nullptr;
    }
}

void D3D12DescriptorSet::Bind() {
}

void D3D12DescriptorSet::Unbind() {
}

void D3D12DescriptorSet::Update(uint32_t offset, RHIResourceView* view) {
    if (m_heap == nullptr || view == nullptr || m_heap->GetType() != DescriptorHeapType::CBV_SRV_UAV) {
        return;
    }

    uint32_t descriptorOffset = 0;
    bool foundBinding = false;
    for (uint32_t i = 0; i < m_bindingCount; ++i) {
        if (m_bindings[i].binding == offset) {
            foundBinding = true;
            break;
        }
        descriptorOffset += m_bindings[i].count;
    }

    if (!foundBinding) {
        return;
    }

    D3D12ResourceView* d3d12View = static_cast<D3D12ResourceView*>(view);
    if (!d3d12View->IsValid()) {
        return;
    }

    CPUDescriptorHandle dstHandle = m_heap->GetCPUDescriptorHandle(m_offset + descriptorOffset);
    D3D12_CPU_DESCRIPTOR_HANDLE dst = { dstHandle.ptr };
    m_heap->GetDevice()->CopyDescriptorsSimple(
        1,
        dst,
        d3d12View->GetCPUHandle(),
        D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);
}

void D3D12DescriptorSet::UpdateSampler(uint32_t offset, RHISampler* sampler) {
    if (m_heap == nullptr || sampler == nullptr || m_heap->GetType() != DescriptorHeapType::Sampler) {
        return;
    }

    uint32_t descriptorOffset = 0;
    bool foundBinding = false;
    for (uint32_t i = 0; i < m_bindingCount; ++i) {
        if (m_bindings[i].binding == offset) {
            foundBinding = true;
            break;
        }
        descriptorOffset += m_bindings[i].count;
    }

    if (!foundBinding) {
        return;
    }

    D3D12Sampler* d3d12Sampler = static_cast<D3D12Sampler*>(sampler);
    CPUDescriptorHandle dstHandle = m_heap->GetCPUDescriptorHandle(m_offset + descriptorOffset);
    D3D12_CPU_DESCRIPTOR_HANDLE dst = { dstHandle.ptr };
    m_heap->GetDevice()->CreateSampler(&d3d12Sampler->GetDesc(), dst);
}

D3D12_GPU_DESCRIPTOR_HANDLE D3D12DescriptorSet::GetGPUHandle(uint32_t index) const {
    if (m_heap == nullptr) {
        return D3D12_GPU_DESCRIPTOR_HANDLE{0};
    }
    GPUDescriptorHandle handle = m_heap->GetGPUDescriptorHandle(m_offset + index);
    return D3D12_GPU_DESCRIPTOR_HANDLE{ handle.ptr };
}

void D3D12DescriptorSet::WriteConstant(uint32_t binding, const void* data, size_t size, size_t offset) {
    size_t requiredSize = offset + size;
    if (m_constantBufferData.size() < requiredSize) {
        m_constantBufferData.resize(requiredSize);
    }
    memcpy(m_constantBufferData.data() + offset, data, size);
    m_constantBufferDirty = true;
}

} // namespace RHI
} // namespace XCEngine