#include "D3D12Resources.h"

namespace XCEngine {
namespace RHI {

D3D12Buffer::D3D12Buffer(D3D12Device* device, ID3D12Resource* resource, const BufferDesc& desc)
    : m_device(device), m_resource(resource), m_desc(desc) {
}

D3D12Buffer::~D3D12Buffer() {
    m_resource.Reset();
}

void* D3D12Buffer::Map() {
    void* data = nullptr;
    D3D12_RANGE range = { 0, (SIZE_T)m_desc.size };
    m_resource->Map(0, &range, &data);
    return data;
}

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

D3D12VertexBuffer::D3D12VertexBuffer(D3D12Buffer* buffer) : m_buffer(buffer) {
}

D3D12VertexBuffer::~D3D12VertexBuffer() {
    delete m_buffer;
}

D3D12IndexBuffer::D3D12IndexBuffer(D3D12Buffer* buffer) : m_buffer(buffer) {
}

D3D12IndexBuffer::~D3D12IndexBuffer() {
    delete m_buffer;
}

D3D12ConstantBuffer::D3D12ConstantBuffer(D3D12Buffer* buffer) : m_buffer(buffer) {
}

D3D12ConstantBuffer::~D3D12ConstantBuffer() {
    delete m_buffer;
}

D3D12Texture2D::D3D12Texture2D(D3D12Device* device, ID3D12Resource* resource, const TextureDesc& desc)
    : m_device(device), m_resource(resource), m_desc(desc) {
}

D3D12Texture2D::~D3D12Texture2D() {
    m_resource.Reset();
}

void D3D12Texture2D::CreateSRV() {
    if (m_hasSRV) return;
    
    D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc = {};
    srvDesc.Format = FormatToDXGIFormat(m_desc.format);
    srvDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
    srvDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
    srvDesc.Texture2D.MipLevels = m_desc.mipLevels;
    
    m_device->GetD3D12Device()->CreateShaderResourceView(m_resource.Get(), &srvDesc, m_srvHandle);
    m_hasSRV = true;
}

void* D3D12Texture2D::GetSRV() {
    return (void*)m_srvHandle.ptr;
}

void D3D12Texture2D::CreateRTV() {
    if (m_hasRTV) return;
    
    D3D12_RENDER_TARGET_VIEW_DESC rtvDesc = {};
    rtvDesc.Format = FormatToDXGIFormat(m_desc.format);
    rtvDesc.ViewDimension = D3D12_RTV_DIMENSION_TEXTURE2D;
    
    m_device->GetD3D12Device()->CreateRenderTargetView(m_resource.Get(), &rtvDesc, m_rtvHandle);
    m_hasRTV = true;
}

void* D3D12Texture2D::GetRTV() {
    return (void*)m_rtvHandle.ptr;
}

void D3D12Texture2D::CreateDSV() {
    if (m_hasDSV) return;
    
    D3D12_DEPTH_STENCIL_VIEW_DESC dsvDesc = {};
    dsvDesc.Format = FormatToDXGIFormat(m_desc.format);
    dsvDesc.ViewDimension = D3D12_DSV_DIMENSION_TEXTURE2D;
    
    m_device->GetD3D12Device()->CreateDepthStencilView(m_resource.Get(), &dsvDesc, m_dsvHandle);
    m_hasDSV = true;
}

void* D3D12Texture2D::GetDSV() {
    return (void*)m_dsvHandle.ptr;
}

bool CreateVertexBuffer(D3D12Device* device, ICommandList* commandList,
    const void* data, uint64_t size, uint32_t stride, IVertexBuffer** outBuffer) {
    
    D3D12_HEAP_PROPERTIES heapProps = {};
    heapProps.Type = D3D12_HEAP_TYPE_DEFAULT;
    
    D3D12_RESOURCE_DESC resourceDesc = {};
    resourceDesc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
    resourceDesc.Width = size;
    resourceDesc.Height = 1;
    resourceDesc.DepthOrArraySize = 1;
    resourceDesc.MipLevels = 1;
    resourceDesc.Format = DXGI_FORMAT_UNKNOWN;
    resourceDesc.SampleDesc.Count = 1;
    resourceDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
    
    ID3D12Resource* bufferResource = nullptr;
    HRESULT hr = device->GetD3D12Device()->CreateCommittedResource(
        &heapProps, D3D12_HEAP_FLAG_NONE, &resourceDesc,
        D3D12_RESOURCE_STATE_COPY_DEST, nullptr, IID_PPV_ARGS(&bufferResource));
    if (FAILED(hr)) return false;
    
    D3D12_RESOURCE_DESC bufferDesc = bufferResource->GetDesc();
    UINT64 memorySize = 0;
    UINT64 rowSize = 0;
    UINT rowCount = 0;
    D3D12_PLACED_SUBRESOURCE_FOOTPRINT footprint;
    device->GetD3D12Device()->GetCopyableFootprints(&bufferDesc, 0, 1, 0, &footprint, &rowCount, &rowSize, &memorySize);
    
    D3D12_HEAP_PROPERTIES uploadHeapProps = {};
    uploadHeapProps.Type = D3D12_HEAP_TYPE_UPLOAD;
    
    ID3D12Resource* uploadResource = nullptr;
    hr = device->GetD3D12Device()->CreateCommittedResource(
        &uploadHeapProps, D3D12_HEAP_FLAG_NONE, &resourceDesc,
        D3D12_RESOURCE_STATE_GENERIC_READ, nullptr, IID_PPV_ARGS(&uploadResource));
    if (FAILED(hr)) {
        bufferResource->Release();
        return false;
    }
    
    UINT8* uploadData = nullptr;
    uploadResource->Map(0, nullptr, (void**)&uploadData);
    
    UINT8* dst = uploadData + footprint.Offset;
    const UINT8* src = (const UINT8*)data;
    for (UINT i = 0; i < rowCount; i++) {
        memcpy(dst + footprint.Footprint.RowPitch * i, src + rowSize * i, (SIZE_T)rowSize);
    }
    uploadResource->Unmap(0, nullptr);
    
    D3D12CommandList* d3dCmdList = static_cast<D3D12CommandList*>(commandList);
    d3dCmdList->GetNative()->CopyBufferRegion(bufferResource, 0, uploadResource, 0, size);
    
    D3D12_RESOURCE_BARRIER barrier = {};
    barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
    barrier.Transition.pResource = bufferResource;
    barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_COPY_DEST;
    barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_VERTEX_AND_CONSTANT_BUFFER;
    barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
    d3dCmdList->GetNative()->ResourceBarrier(1, &barrier);
    
    uploadResource->Release();
    
    BufferDesc bufferDesc2 = {};
    bufferDesc2.size = size;
    bufferDesc2.stride = stride;
    
    D3D12Buffer* buffer = new D3D12Buffer(device, bufferResource, bufferDesc2);
    *outBuffer = new D3D12VertexBuffer(buffer);
    
    return true;
}

bool CreateIndexBuffer(D3D12Device* device, ICommandList* commandList,
    const void* data, uint64_t size, IIndexBuffer** outBuffer) {
    
    D3D12_HEAP_PROPERTIES heapProps = {};
    heapProps.Type = D3D12_HEAP_TYPE_DEFAULT;
    
    D3D12_RESOURCE_DESC resourceDesc = {};
    resourceDesc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
    resourceDesc.Width = size;
    resourceDesc.Height = 1;
    resourceDesc.DepthOrArraySize = 1;
    resourceDesc.MipLevels = 1;
    resourceDesc.Format = DXGI_FORMAT_UNKNOWN;
    resourceDesc.SampleDesc.Count = 1;
    resourceDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
    
    ID3D12Resource* bufferResource = nullptr;
    HRESULT hr = device->GetD3D12Device()->CreateCommittedResource(
        &heapProps, D3D12_HEAP_FLAG_NONE, &resourceDesc,
        D3D12_RESOURCE_STATE_COPY_DEST, nullptr, IID_PPV_ARGS(&bufferResource));
    if (FAILED(hr)) return false;
    
    D3D12_RESOURCE_DESC bufferDesc = bufferResource->GetDesc();
    UINT64 memorySize = 0;
    UINT64 rowSize = 0;
    UINT rowCount = 0;
    D3D12_PLACED_SUBRESOURCE_FOOTPRINT footprint;
    device->GetD3D12Device()->GetCopyableFootprints(&bufferDesc, 0, 1, 0, &footprint, &rowCount, &rowSize, &memorySize);
    
    D3D12_HEAP_PROPERTIES uploadHeapProps = {};
    uploadHeapProps.Type = D3D12_HEAP_TYPE_UPLOAD;
    
    ID3D12Resource* uploadResource = nullptr;
    hr = device->GetD3D12Device()->CreateCommittedResource(
        &uploadHeapProps, D3D12_HEAP_FLAG_NONE, &resourceDesc,
        D3D12_RESOURCE_STATE_GENERIC_READ, nullptr, IID_PPV_ARGS(&uploadResource));
    if (FAILED(hr)) {
        bufferResource->Release();
        return false;
    }
    
    UINT8* uploadData = nullptr;
    uploadResource->Map(0, nullptr, (void**)&uploadData);
    
    memcpy(uploadData + footprint.Offset, data, (SIZE_T)size);
    uploadResource->Unmap(0, nullptr);
    
    D3D12CommandList* d3dCmdList = static_cast<D3D12CommandList*>(commandList);
    d3dCmdList->GetNative()->CopyBufferRegion(bufferResource, 0, uploadResource, 0, size);
    
    D3D12_RESOURCE_BARRIER barrier = {};
    barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
    barrier.Transition.pResource = bufferResource;
    barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_COPY_DEST;
    barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_INDEX_BUFFER;
    barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
    d3dCmdList->GetNative()->ResourceBarrier(1, &barrier);
    
    uploadResource->Release();
    
    BufferDesc bufferDesc2 = {};
    bufferDesc2.size = size;
    bufferDesc2.stride = 4;
    
    D3D12Buffer* buffer = new D3D12Buffer(device, bufferResource, bufferDesc2);
    *outBuffer = new D3D12IndexBuffer(buffer);
    
    return true;
}

bool CreateConstantBuffer(D3D12Device* device, uint64_t size, IConstantBuffer** outBuffer) {
    D3D12_HEAP_PROPERTIES heapProps = {};
    heapProps.Type = D3D12_HEAP_TYPE_UPLOAD;
    
    D3D12_RESOURCE_DESC resourceDesc = {};
    resourceDesc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
    resourceDesc.Width = size;
    resourceDesc.Height = 1;
    resourceDesc.DepthOrArraySize = 1;
    resourceDesc.MipLevels = 1;
    resourceDesc.Format = DXGI_FORMAT_UNKNOWN;
    resourceDesc.SampleDesc.Count = 1;
    resourceDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
    
    ID3D12Resource* bufferResource = nullptr;
    HRESULT hr = device->GetD3D12Device()->CreateCommittedResource(
        &heapProps, D3D12_HEAP_FLAG_NONE, &resourceDesc,
        D3D12_RESOURCE_STATE_GENERIC_READ, nullptr, IID_PPV_ARGS(&bufferResource));
    if (FAILED(hr)) return false;
    
    BufferDesc bufferDesc = {};
    bufferDesc.size = size;
    bufferDesc.stride = (uint32_t)size;
    bufferDesc.cpuAccessible = true;
    
    D3D12Buffer* buffer = new D3D12Buffer(device, bufferResource, bufferDesc);
    *outBuffer = new D3D12ConstantBuffer(buffer);
    
    return true;
}

bool CreateTexture2D(D3D12Device* device, ICommandList* commandList,
    const void* pixels, uint32_t width, uint32_t height, Format format, ITexture2D** outTexture) {
    
    D3D12_HEAP_PROPERTIES heapProps = {};
    heapProps.Type = D3D12_HEAP_TYPE_DEFAULT;
    
    D3D12_RESOURCE_DESC resourceDesc = {};
    resourceDesc.Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D;
    resourceDesc.Width = width;
    resourceDesc.Height = height;
    resourceDesc.DepthOrArraySize = 1;
    resourceDesc.MipLevels = 1;
    resourceDesc.Format = FormatToDXGIFormat(format);
    resourceDesc.SampleDesc.Count = 1;
    resourceDesc.Layout = D3D12_TEXTURE_LAYOUT_UNKNOWN;
    
    ID3D12Resource* textureResource = nullptr;
    HRESULT hr = device->GetD3D12Device()->CreateCommittedResource(
        &heapProps, D3D12_HEAP_FLAG_NONE, &resourceDesc,
        D3D12_RESOURCE_STATE_COPY_DEST, nullptr, IID_PPV_ARGS(&textureResource));
    if (FAILED(hr)) return false;
    
    D3D12_RESOURCE_DESC texDesc = textureResource->GetDesc();
    UINT64 memorySize = 0;
    UINT64 rowSize = 0;
    UINT rowCount = 0;
    D3D12_PLACED_SUBRESOURCE_FOOTPRINT footprint;
    device->GetD3D12Device()->GetCopyableFootprints(&texDesc, 0, 1, 0, &footprint, &rowCount, &rowSize, &memorySize);
    
    D3D12_HEAP_PROPERTIES uploadHeapProps = {};
    uploadHeapProps.Type = D3D12_HEAP_TYPE_UPLOAD;
    
    D3D12_RESOURCE_DESC uploadDesc = {};
    uploadDesc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
    uploadDesc.Width = memorySize;
    uploadDesc.Height = 1;
    uploadDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
    
    ID3D12Resource* uploadResource = nullptr;
    hr = device->GetD3D12Device()->CreateCommittedResource(
        &uploadHeapProps, D3D12_HEAP_FLAG_NONE, &uploadDesc,
        D3D12_RESOURCE_STATE_GENERIC_READ, nullptr, IID_PPV_ARGS(&uploadResource));
    if (FAILED(hr)) {
        textureResource->Release();
        return false;
    }
    
    UINT8* uploadData = nullptr;
    uploadResource->Map(0, nullptr, (void**)&uploadData);
    
    UINT8* dst = uploadData + footprint.Offset;
    const UINT8* src = (const UINT8*)pixels;
    for (UINT i = 0; i < rowCount; i++) {
        memcpy(dst + footprint.Footprint.RowPitch * i, src + rowSize * i, (SIZE_T)rowSize);
    }
    uploadResource->Unmap(0, nullptr);
    
    D3D12CommandList* d3dCmdList = static_cast<D3D12CommandList*>(commandList);
    
    D3D12_TEXTURE_COPY_LOCATION dstLoc = {};
    dstLoc.pResource = textureResource;
    dstLoc.Type = D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX;
    dstLoc.SubresourceIndex = 0;
    
    D3D12_TEXTURE_COPY_LOCATION srcLoc = {};
    srcLoc.pResource = uploadResource;
    srcLoc.Type = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT;
    srcLoc.PlacedFootprint = footprint;
    
    d3dCmdList->GetNative()->CopyTextureRegion(&dstLoc, 0, 0, 0, &srcLoc, nullptr);
    
    D3D12_RESOURCE_BARRIER barrier = {};
    barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
    barrier.Transition.pResource = textureResource;
    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;
    d3dCmdList->GetNative()->ResourceBarrier(1, &barrier);
    
    uploadResource->Release();
    
    TextureDesc texDesc2 = {};
    texDesc2.width = width;
    texDesc2.height = height;
    texDesc2.format = format;
    
    D3D12Texture2D* texture = new D3D12Texture2D(device, textureResource, texDesc2);
    *outTexture = texture;
    
    return true;
}

void UpdateConstantBuffer(IConstantBuffer* buffer, const void* data, uint64_t size) {
    D3D12ConstantBuffer* d3dBuffer = static_cast<D3D12ConstantBuffer*>(buffer);
    void* mapped = d3dBuffer->Map();
    if (mapped) {
        memcpy(mapped, data, (SIZE_T)size);
        d3dBuffer->Unmap();
    }
}

} // namespace RHI
} // namespace XCEngine