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

namespace XCEngine {
namespace RHI {

D3D12Buffer::D3D12Buffer() {
}

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

bool D3D12Buffer::Initialize(ID3D12Device* device, uint64_t size, D3D12_RESOURCE_STATES initialState, D3D12_HEAP_TYPE heapType) {
    D3D12_HEAP_PROPERTIES heapProperties = {};
    heapProperties.Type = heapType;
    heapProperties.CPUPageProperty = D3D12_CPU_PAGE_PROPERTY_UNKNOWN;
    heapProperties.MemoryPoolPreference = D3D12_MEMORY_POOL_UNKNOWN;
    heapProperties.CreationNodeMask = 0;
    heapProperties.VisibleNodeMask = 0;

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

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

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

    return true;
}

bool D3D12Buffer::InitializeFromExisting(ID3D12Resource* resource) {
    m_resource = resource;
    return true;
}

bool D3D12Buffer::InitializeWithData(ID3D12Device* device, ID3D12GraphicsCommandList* commandList, 
    const void* data, uint64_t size, D3D12_RESOURCE_STATES finalState) {
    
    D3D12_HEAP_PROPERTIES heapProperties = {};
    heapProperties.Type = D3D12_HEAP_TYPE_DEFAULT;
    
    D3D12_RESOURCE_DESC bufferDesc = {};
    bufferDesc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
    bufferDesc.Alignment = 0;
    bufferDesc.Width = size;
    bufferDesc.Height = 1;
    bufferDesc.DepthOrArraySize = 1;
    bufferDesc.MipLevels = 1;
    bufferDesc.Format = DXGI_FORMAT_UNKNOWN;
    bufferDesc.SampleDesc.Count = 1;
    bufferDesc.SampleDesc.Quality = 0;
    bufferDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
    bufferDesc.Flags = D3D12_RESOURCE_FLAG_NONE;

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

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

    bufferDesc = m_resource->GetDesc();
    UINT64 memorySizeUsed = 0;
    UINT64 rowSizeInBytes = 0;
    UINT rowUsed = 0;
    D3D12_PLACED_SUBRESOURCE_FOOTPRINT subresourceFootprint;
    device->GetCopyableFootprints(&bufferDesc, 0, 1, 0,
        &subresourceFootprint, &rowUsed, &rowSizeInBytes, &memorySizeUsed);

    ID3D12Resource* tempBufferObject = nullptr;
    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 = size;
    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)
    );

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

    BYTE* pData;
    tempBufferObject->Map(0, nullptr, reinterpret_cast<void**>(&pData));
    BYTE* pDstTempBuffer = reinterpret_cast<BYTE*>(pData + subresourceFootprint.Offset);
    const BYTE* pSrcData = reinterpret_cast<const BYTE*>(data);
    memcpy(pDstTempBuffer, pSrcData, size);
    tempBufferObject->Unmap(0, nullptr);

    commandList->CopyBufferRegion(m_resource.Get(), 0, tempBufferObject, 0, size);

    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 = finalState;
    barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
    commandList->ResourceBarrier(1, &barrier);

    return true;
}

void D3D12Buffer::Shutdown() {
    m_resource.Reset();
}

void D3D12Buffer::UpdateData(const void* data, uint64_t size) {
    D3D12_RANGE range = { 0, static_cast<SIZE_T>(size) };
    unsigned char* pBuffer = nullptr;
    m_resource->Map(0, &range, reinterpret_cast<void**>(&pBuffer));
    memcpy(pBuffer, data, static_cast<size_t>(size));
    m_resource->Unmap(0, nullptr);
}

void* D3D12Buffer::Map() {
    D3D12_RANGE readRange = { 0, 0 };
    void* data = nullptr;
    m_resource->Map(0, &readRange, &data);
    return data;
}

void D3D12Buffer::Unmap() {
    m_resource->Unmap(0, nullptr);
}

void D3D12Buffer::SetData(const void* data, size_t size, size_t offset) {
    D3D12_RANGE writeRange = { offset, offset + size };
    unsigned char* pBuffer = nullptr;
    m_resource->Map(0, &writeRange, reinterpret_cast<void**>(&pBuffer));
    memcpy(pBuffer + offset, data, size);
    m_resource->Unmap(0, nullptr);
}

} // namespace RHI
} // namespace XCEngine