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

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_ownsResource = true;
    return true;
}

bool D3D12Texture::InitializeFromExisting(ID3D12Resource* resource, bool ownsResource) {
    m_resource = resource;
    m_ownsResource = ownsResource;
    return true;
}

bool D3D12Texture::InitializeFromData(ID3D12Device* device, ID3D12GraphicsCommandList* commandList, 
    const void* pixelData, uint32_t width, uint32_t height, DXGI_FORMAT format, uint32_t rowPitch,
    ComPtr<ID3D12Resource>* uploadBuffer) {
    
    D3D12_RESOURCE_DESC textureDesc = {};
    textureDesc.Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D;
    textureDesc.Alignment = 0;
    textureDesc.Width = width;
    textureDesc.Height = height;
    textureDesc.DepthOrArraySize = 1;
    textureDesc.MipLevels = 1;
    textureDesc.Format = format;
    textureDesc.SampleDesc.Count = 1;
    textureDesc.SampleDesc.Quality = 0;
    textureDesc.Layout = D3D12_TEXTURE_LAYOUT_UNKNOWN;
    textureDesc.Flags = D3D12_RESOURCE_FLAG_NONE;

    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;

    textureDesc = m_resource->GetDesc();
    UINT64 memorySizeUsed = 0;
    UINT64 rowSizeInBytes = 0;
    UINT rowUsed = 0;
    D3D12_PLACED_SUBRESOURCE_FOOTPRINT subresourceFootprint;
    device->GetCopyableFootprints(&textureDesc, 0, 1, 0,
        &subresourceFootprint, &rowUsed, &rowSizeInBytes, &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;
    tempBufferObject->Map(0, nullptr, reinterpret_cast<void**>(&pData));
    BYTE* pDstTempBuffer = reinterpret_cast<BYTE*>(pData + subresourceFootprint.Offset);
    const BYTE* pSrcData = reinterpret_cast<const BYTE*>(pixelData);
    const UINT sourceRowPitch = rowPitch > 0 ? rowPitch : static_cast<UINT>(rowSizeInBytes);
    for (UINT i = 0; i < rowUsed; i++) {
        memcpy(pDstTempBuffer + subresourceFootprint.Footprint.RowPitch * i, pSrcData + sourceRowPitch * i, rowSizeInBytes);
    }
    tempBufferObject->Unmap(0, nullptr);

    D3D12_TEXTURE_COPY_LOCATION dst = {};
    dst.pResource = m_resource.Get();
    dst.Type = D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX;
    dst.SubresourceIndex = 0;

    D3D12_TEXTURE_COPY_LOCATION src = {};
    src.pResource = tempBufferObject.Get();
    src.Type = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT;
    src.PlacedFootprint = subresourceFootprint;
    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) {
    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 = format;
    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)
    );

    return SUCCEEDED(hResult);
}

void D3D12Texture::Shutdown() {
    if (m_ownsResource) {
        m_resource.Reset();
    }
}

} // namespace RHI
} // namespace XCEngine