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

#include <algorithm>
#include <vector>

namespace XCEngine {
namespace RHI {

D3D12Texture::D3D12Texture() {
}

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

bool D3D12Texture::Initialize(ID3D12Device* device, const D3D12_RESOURCE_DESC& desc, D3D12_RESOURCE_STATES initialState) {
    D3D12_HEAP_PROPERTIES heapProperties = {};
    heapProperties.Type = D3D12_HEAP_TYPE_DEFAULT;
    heapProperties.CPUPageProperty = D3D12_CPU_PAGE_PROPERTY_UNKNOWN;
    heapProperties.MemoryPoolPreference = D3D12_MEMORY_POOL_UNKNOWN;
    heapProperties.CreationNodeMask = 0;
    heapProperties.VisibleNodeMask = 0;

    D3D12_CLEAR_VALUE* pOptimizedClearValue = nullptr;

    HRESULT hResult = device->CreateCommittedResource(
        &heapProperties,
        D3D12_HEAP_FLAG_NONE,
        &desc,
        initialState,
        pOptimizedClearValue,
        IID_PPV_ARGS(&m_resource)
    );

    if (FAILED(hResult)) {
        return false;
    }

    m_format = FromD3D12(desc.Format);
    switch (desc.Dimension) {
        case D3D12_RESOURCE_DIMENSION_TEXTURE1D:
            m_textureType = TextureType::Texture1D;
            break;
        case D3D12_RESOURCE_DIMENSION_TEXTURE3D:
            m_textureType = TextureType::Texture3D;
            break;
        case D3D12_RESOURCE_DIMENSION_TEXTURE2D:
        default:
            m_textureType = TextureType::Texture2D;
            break;
    }
    m_ownsResource = true;
    return true;
}

bool D3D12Texture::InitializeFromExisting(ID3D12Resource* resource, bool ownsResource) {
    m_resource = resource;
    if (m_resource) {
        const D3D12_RESOURCE_DESC desc = m_resource->GetDesc();
        m_format = FromD3D12(desc.Format);
        switch (desc.Dimension) {
            case D3D12_RESOURCE_DIMENSION_TEXTURE1D:
                m_textureType = TextureType::Texture1D;
                break;
            case D3D12_RESOURCE_DIMENSION_TEXTURE3D:
                m_textureType = TextureType::Texture3D;
                break;
            case D3D12_RESOURCE_DIMENSION_TEXTURE2D:
            default:
                m_textureType = TextureType::Texture2D;
                break;
        }
    }
    m_ownsResource = ownsResource;
    return true;
}

bool D3D12Texture::InitializeFromData(ID3D12Device* device, ID3D12GraphicsCommandList* commandList,
    const D3D12_RESOURCE_DESC& textureDesc, TextureType textureType, const void* pixelData, size_t pixelDataSize,
    uint32_t rowPitch, ComPtr<ID3D12Resource>* uploadBuffer) {
    if (device == nullptr || commandList == nullptr || pixelData == nullptr || pixelDataSize == 0) {
        return false;
    }

    D3D12_HEAP_PROPERTIES heapProperties = {};
    heapProperties.Type = D3D12_HEAP_TYPE_DEFAULT;

    HRESULT hResult = device->CreateCommittedResource(
        &heapProperties,
        D3D12_HEAP_FLAG_NONE,
        &textureDesc,
        D3D12_RESOURCE_STATE_COPY_DEST,
        nullptr,
        IID_PPV_ARGS(&m_resource)
    );

    if (FAILED(hResult)) {
        return false;
    }
    m_ownsResource = true;
    m_format = FromD3D12(textureDesc.Format);
    m_textureType = textureType;

    const UINT subresourceCount = textureDesc.MipLevels *
        (textureDesc.Dimension == D3D12_RESOURCE_DIMENSION_TEXTURE3D ? 1u : textureDesc.DepthOrArraySize);
    std::vector<D3D12_PLACED_SUBRESOURCE_FOOTPRINT> subresourceFootprints(subresourceCount);
    std::vector<UINT> subresourceRowCounts(subresourceCount);
    std::vector<UINT64> subresourceRowSizes(subresourceCount);
    UINT64 memorySizeUsed = 0;
    device->GetCopyableFootprints(
        &textureDesc,
        0,
        subresourceCount,
        0,
        subresourceFootprints.data(),
        subresourceRowCounts.data(),
        subresourceRowSizes.data(),
        &memorySizeUsed);

    ComPtr<ID3D12Resource> tempBufferObject;
    D3D12_HEAP_PROPERTIES d3dTempHeapProperties = {};
    d3dTempHeapProperties.Type = D3D12_HEAP_TYPE_UPLOAD;

    D3D12_RESOURCE_DESC d3d12TempResourceDesc = {};
    d3d12TempResourceDesc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
    d3d12TempResourceDesc.Alignment = 0;
    d3d12TempResourceDesc.Width = memorySizeUsed;
    d3d12TempResourceDesc.Height = 1;
    d3d12TempResourceDesc.DepthOrArraySize = 1;
    d3d12TempResourceDesc.MipLevels = 1;
    d3d12TempResourceDesc.Format = DXGI_FORMAT_UNKNOWN;
    d3d12TempResourceDesc.SampleDesc.Count = 1;
    d3d12TempResourceDesc.SampleDesc.Quality = 0;
    d3d12TempResourceDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
    d3d12TempResourceDesc.Flags = D3D12_RESOURCE_FLAG_NONE;

    hResult = device->CreateCommittedResource(
        &d3dTempHeapProperties,
        D3D12_HEAP_FLAG_NONE,
        &d3d12TempResourceDesc,
        D3D12_RESOURCE_STATE_GENERIC_READ,
        nullptr,
        IID_PPV_ARGS(tempBufferObject.GetAddressOf())
    );

    if (FAILED(hResult)) {
        return false;
    }

    BYTE* pData = nullptr;
    tempBufferObject->Map(0, nullptr, reinterpret_cast<void**>(&pData));
    const BYTE* pSrcData = reinterpret_cast<const BYTE*>(pixelData);
    const UINT sourceRowPitch = rowPitch > 0 ? rowPitch : static_cast<UINT>(subresourceRowSizes[0]);

    size_t sourceOffset = 0;
    for (UINT subresourceIndex = 0; subresourceIndex < subresourceCount; ++subresourceIndex) {
        const D3D12_PLACED_SUBRESOURCE_FOOTPRINT& footprint = subresourceFootprints[subresourceIndex];
        const UINT rowCount = subresourceRowCounts[subresourceIndex];
        const UINT64 rowSizeInBytes = subresourceRowSizes[subresourceIndex];
        BYTE* pDstTempBuffer = pData + footprint.Offset;

        for (UINT depthSlice = 0; depthSlice < footprint.Footprint.Depth; ++depthSlice) {
            for (UINT rowIndex = 0; rowIndex < rowCount; ++rowIndex) {
                const size_t srcRowOffset =
                    sourceOffset +
                    static_cast<size_t>(depthSlice) * static_cast<size_t>(sourceRowPitch) * rowCount +
                    static_cast<size_t>(rowIndex) * sourceRowPitch;
                memcpy(
                    pDstTempBuffer +
                        static_cast<size_t>(depthSlice) * footprint.Footprint.RowPitch * rowCount +
                        static_cast<size_t>(rowIndex) * footprint.Footprint.RowPitch,
                    pSrcData + srcRowOffset,
                    static_cast<size_t>(std::min<UINT64>(rowSizeInBytes, sourceRowPitch)));
            }
        }

        sourceOffset += static_cast<size_t>(sourceRowPitch) * rowCount * footprint.Footprint.Depth;
    }
    tempBufferObject->Unmap(0, nullptr);

    for (UINT subresourceIndex = 0; subresourceIndex < subresourceCount; ++subresourceIndex) {
        D3D12_TEXTURE_COPY_LOCATION dst = {};
        dst.pResource = m_resource.Get();
        dst.Type = D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX;
        dst.SubresourceIndex = subresourceIndex;

        D3D12_TEXTURE_COPY_LOCATION src = {};
        src.pResource = tempBufferObject.Get();
        src.Type = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT;
        src.PlacedFootprint = subresourceFootprints[subresourceIndex];
        commandList->CopyTextureRegion(&dst, 0, 0, 0, &src, nullptr);
    }

    D3D12_RESOURCE_BARRIER barrier = {};
    barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
    barrier.Flags = D3D12_RESOURCE_BARRIER_FLAG_NONE;
    barrier.Transition.pResource = m_resource.Get();
    barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_COPY_DEST;
    barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE;
    barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
    commandList->ResourceBarrier(1, &barrier);
    m_state = ResourceStates::PixelShaderResource;

    if (uploadBuffer != nullptr) {
        *uploadBuffer = tempBufferObject;
    }

    return true;
}

bool D3D12Texture::InitializeDepthStencil(ID3D12Device* device, uint32_t width, uint32_t height, DXGI_FORMAT format) {
    DXGI_FORMAT resourceFormat = format;
    switch (format) {
        case DXGI_FORMAT_D16_UNORM:
            resourceFormat = DXGI_FORMAT_R16_TYPELESS;
            break;
        case DXGI_FORMAT_D24_UNORM_S8_UINT:
            resourceFormat = DXGI_FORMAT_R24G8_TYPELESS;
            break;
        case DXGI_FORMAT_D32_FLOAT:
            resourceFormat = DXGI_FORMAT_R32_TYPELESS;
            break;
        default:
            break;
    }

    D3D12_RESOURCE_DESC desc = {};
    desc.Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D;
    desc.Alignment = 0;
    desc.Width = width;
    desc.Height = height;
    desc.DepthOrArraySize = 1;
    desc.MipLevels = 1;
    desc.Format = resourceFormat;
    desc.SampleDesc.Count = 1;
    desc.SampleDesc.Quality = 0;
    desc.Layout = D3D12_TEXTURE_LAYOUT_UNKNOWN;
    desc.Flags = D3D12_RESOURCE_FLAG_ALLOW_DEPTH_STENCIL;

    D3D12_HEAP_PROPERTIES heapProperties = {};
    heapProperties.Type = D3D12_HEAP_TYPE_DEFAULT;

    D3D12_CLEAR_VALUE clearValue = {};
    clearValue.Format = format;
    clearValue.DepthStencil.Depth = 1.0f;
    clearValue.DepthStencil.Stencil = 0;

    HRESULT hResult = device->CreateCommittedResource(
        &heapProperties,
        D3D12_HEAP_FLAG_NONE,
        &desc,
        D3D12_RESOURCE_STATE_DEPTH_WRITE,
        &clearValue,
        IID_PPV_ARGS(&m_resource)
    );

    if (SUCCEEDED(hResult)) {
        m_format = FromD3D12(format);
        m_textureType = TextureType::Texture2D;
        m_ownsResource = true;
        return true;
    }

    return false;
}

void D3D12Texture::Shutdown() {
    m_resource.Reset();
    m_state = ResourceStates::Common;
    m_format = Format::Unknown;
    m_textureType = TextureType::Texture2D;
    m_ownsResource = false;
}

} // namespace RHI
} // namespace XCEngine