#include "D3D12ShaderResourceDescriptorAllocator.h"

#include <XCEngine/RHI/RHITypes.h>

namespace XCEngine::UI::Editor::Host {

namespace {

using ::XCEngine::RHI::D3D12DescriptorHeap;
using ::XCEngine::RHI::D3D12Device;
using ::XCEngine::RHI::D3D12Texture;
using ::XCEngine::RHI::DescriptorHeapType;
using ::XCEngine::RHI::DescriptorPoolDesc;
using ::XCEngine::RHI::RHIDevice;
using ::XCEngine::RHI::RHITexture;

} // namespace

bool D3D12ShaderResourceDescriptorAllocator::Initialize(RHIDevice& device, UINT descriptorCount) {
    Shutdown();

    DescriptorPoolDesc descriptorPoolDesc = {};
    descriptorPoolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    descriptorPoolDesc.descriptorCount = descriptorCount;
    descriptorPoolDesc.shaderVisible = true;
    m_descriptorPool = device.CreateDescriptorPool(descriptorPoolDesc);
    m_descriptorHeap = dynamic_cast<D3D12DescriptorHeap*>(m_descriptorPool);
    if (m_descriptorPool == nullptr || m_descriptorHeap == nullptr) {
        Shutdown();
        return false;
    }

    m_device = &device;
    m_descriptorSize = m_descriptorHeap->GetDescriptorSize();
    m_descriptorCount = descriptorCount;
    m_descriptorUsage.assign(descriptorCount, false);
    return true;
}

void D3D12ShaderResourceDescriptorAllocator::Shutdown() {
    if (m_descriptorPool != nullptr) {
        m_descriptorPool->Shutdown();
        delete m_descriptorPool;
        m_descriptorPool = nullptr;
    }

    m_device = nullptr;
    m_descriptorHeap = nullptr;
    m_descriptorUsage.clear();
    m_descriptorSize = 0u;
    m_descriptorCount = 0u;
}

bool D3D12ShaderResourceDescriptorAllocator::IsInitialized() const {
    return m_device != nullptr &&
        m_descriptorHeap != nullptr &&
        m_descriptorSize > 0u &&
        !m_descriptorUsage.empty();
}

ID3D12DescriptorHeap* D3D12ShaderResourceDescriptorAllocator::GetDescriptorHeap() const {
    return m_descriptorHeap != nullptr ? m_descriptorHeap->GetDescriptorHeap() : nullptr;
}

UINT D3D12ShaderResourceDescriptorAllocator::GetDescriptorSize() const {
    return m_descriptorSize;
}

UINT D3D12ShaderResourceDescriptorAllocator::GetDescriptorCount() const {
    return m_descriptorCount;
}

void D3D12ShaderResourceDescriptorAllocator::Allocate(
    D3D12_CPU_DESCRIPTOR_HANDLE* outCpuHandle,
    D3D12_GPU_DESCRIPTOR_HANDLE* outGpuHandle) {
    AllocateInternal(outCpuHandle, outGpuHandle);
}

bool D3D12ShaderResourceDescriptorAllocator::CreateTextureDescriptor(
    RHITexture* texture,
    D3D12_CPU_DESCRIPTOR_HANDLE* outCpuHandle,
    D3D12_GPU_DESCRIPTOR_HANDLE* outGpuHandle) {
    if (m_device == nullptr ||
        texture == nullptr ||
        outCpuHandle == nullptr ||
        outGpuHandle == nullptr) {
        return false;
    }

    *outCpuHandle = {};
    *outGpuHandle = {};

    auto* nativeDevice = dynamic_cast<D3D12Device*>(m_device);
    auto* nativeTexture = dynamic_cast<D3D12Texture*>(texture);
    if (nativeDevice == nullptr ||
        nativeTexture == nullptr ||
        nativeTexture->GetResource() == nullptr) {
        return false;
    }

    AllocateInternal(outCpuHandle, outGpuHandle);
    if (outCpuHandle->ptr == 0 || outGpuHandle->ptr == 0) {
        *outCpuHandle = {};
        *outGpuHandle = {};
        return false;
    }

    D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc = {};
    srvDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
    srvDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    srvDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
    srvDesc.Texture2D.MipLevels = 1;
    nativeDevice->GetDevice()->CreateShaderResourceView(
        nativeTexture->GetResource(),
        &srvDesc,
        *outCpuHandle);
    return true;
}

void D3D12ShaderResourceDescriptorAllocator::Free(
    D3D12_CPU_DESCRIPTOR_HANDLE cpuHandle,
    D3D12_GPU_DESCRIPTOR_HANDLE) {
    if (m_descriptorHeap == nullptr ||
        m_descriptorSize == 0u ||
        cpuHandle.ptr < m_descriptorHeap->GetCPUDescriptorHandleForHeapStart().ptr) {
        return;
    }

    const SIZE_T offset =
        cpuHandle.ptr - m_descriptorHeap->GetCPUDescriptorHandleForHeapStart().ptr;
    const std::size_t index =
        static_cast<std::size_t>(offset / m_descriptorSize);
    if (index < m_descriptorUsage.size()) {
        m_descriptorUsage[index] = false;
    }
}

void D3D12ShaderResourceDescriptorAllocator::AllocateInternal(
    D3D12_CPU_DESCRIPTOR_HANDLE* outCpuHandle,
    D3D12_GPU_DESCRIPTOR_HANDLE* outGpuHandle) {
    if (outCpuHandle == nullptr || outGpuHandle == nullptr) {
        return;
    }

    *outCpuHandle = {};
    *outGpuHandle = {};
    if (m_descriptorHeap == nullptr ||
        m_descriptorSize == 0u ||
        m_descriptorUsage.empty()) {
        return;
    }

    for (std::size_t index = 0; index < m_descriptorUsage.size(); ++index) {
        if (m_descriptorUsage[index]) {
            continue;
        }

        m_descriptorUsage[index] = true;
        outCpuHandle->ptr =
            m_descriptorHeap->GetCPUDescriptorHandleForHeapStart().ptr +
            static_cast<SIZE_T>(index) * m_descriptorSize;
        outGpuHandle->ptr =
            m_descriptorHeap->GetGPUDescriptorHandleForHeapStart().ptr +
            static_cast<UINT64>(index) * m_descriptorSize;
        return;
    }
}

} // namespace XCEngine::UI::Editor::Host