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

namespace XCEngine {
namespace RHI {

namespace {

uint64_t AlignConstantBufferSize(size_t size) {
    const uint64_t minSize = size > 0 ? static_cast<uint64_t>(size) : 1ull;
    return (minSize + 255ull) & ~255ull;
}

} // namespace

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;
    m_constantBuffer.reset();
    m_constantBufferCapacity = 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) {
    (void)binding;
    size_t requiredSize = offset + size;
    if (m_constantBufferData.size() < requiredSize) {
        m_constantBufferData.resize(requiredSize);
    }
    memcpy(m_constantBufferData.data() + offset, data, size);
    m_constantBufferDirty = true;
}

bool D3D12DescriptorSet::HasBindingType(DescriptorType type) const {
    for (uint32_t i = 0; i < m_bindingCount; ++i) {
        if (static_cast<DescriptorType>(m_bindings[i].type) == type) {
            return true;
        }
    }
    return false;
}

uint32_t D3D12DescriptorSet::GetFirstBindingOfType(DescriptorType type) const {
    for (uint32_t i = 0; i < m_bindingCount; ++i) {
        if (static_cast<DescriptorType>(m_bindings[i].type) == type) {
            return m_bindings[i].binding;
        }
    }
    return UINT32_MAX;
}

bool D3D12DescriptorSet::UploadConstantBuffer() {
    if (!HasBindingType(DescriptorType::CBV) || m_heap == nullptr || m_heap->GetDevice() == nullptr) {
        return false;
    }

    const uint64_t alignedSize = AlignConstantBufferSize(m_constantBufferData.size());
    if (!m_constantBuffer || m_constantBufferCapacity < alignedSize) {
        auto constantBuffer = std::make_unique<D3D12Buffer>();
        if (!constantBuffer->Initialize(
                m_heap->GetDevice(),
                alignedSize,
                D3D12_RESOURCE_STATE_GENERIC_READ,
                D3D12_HEAP_TYPE_UPLOAD)) {
            return false;
        }

        constantBuffer->SetBufferType(BufferType::Constant);
        constantBuffer->SetStride(static_cast<uint32_t>(alignedSize));
        m_constantBuffer = std::move(constantBuffer);
        m_constantBufferCapacity = alignedSize;
        m_constantBufferDirty = true;
    }

    if (m_constantBufferDirty && !m_constantBufferData.empty()) {
        m_constantBuffer->SetData(m_constantBufferData.data(), m_constantBufferData.size());
        m_constantBufferDirty = false;
    }

    return m_constantBuffer != nullptr;
}

D3D12_GPU_VIRTUAL_ADDRESS D3D12DescriptorSet::GetConstantBufferGPUAddress() const {
    return m_constantBuffer ? m_constantBuffer->GetGPUVirtualAddress() : 0;
}

} // namespace RHI
} // namespace XCEngine