/*
** Copyright (c) 2021 LunarG, Inc.
** Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved.
**
** Permission is hereby granted, free of charge, to any person obtaining a
** copy of this software and associated documentation files (the "Software"),
** to deal in the Software without restriction, including without limitation
** the rights to use, copy, modify, merge, publish, distribute, sublicense,
** and/or sell copies of the Software, and to permit persons to whom the
** Software is furnished to do so, subject to the following conditions:
**
** The above copyright notice and this permission notice shall be included in
** all copies or substantial portions of the Software.
**
** THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
** IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
** FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
** AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
** LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
** FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
** DEALINGS IN THE SOFTWARE.
*/

#include "graphics/dx12_resource_data_util.h"

#include "graphics/dx12_util.h"
#include "util/logging.h"

#include <inttypes.h>
#include <algorithm>

GFXRECON_BEGIN_NAMESPACE(gfxrecon)
GFXRECON_BEGIN_NAMESPACE(graphics)

HRESULT
Dx12ResourceDataUtil::MapSubresourceAndReadData(ID3D12Resource* resource, UINT subresource, size_t size, uint8_t* data)
{
    uint8_t* subresource_data;
    HRESULT  result = dx12::MapSubresource(resource, subresource, nullptr, subresource_data);
    if (SUCCEEDED(result))
    {
        util::platform::MemoryCopy(data, size, subresource_data, size);
        resource->Unmap(subresource, &dx12::kZeroRange);
    }
    return result;
}

HRESULT Dx12ResourceDataUtil::MapSubresourceAndWriteData(ID3D12Resource* resource,
                                                         UINT            subresource,
                                                         size_t          size,
                                                         const uint8_t*  data)
{
    uint8_t* subresource_data;
    HRESULT  result = dx12::MapSubresource(resource, subresource, &dx12::kZeroRange, subresource_data);
    if (SUCCEEDED(result))
    {
        util::platform::MemoryCopy(subresource_data, size, data, size);
        resource->Unmap(subresource, nullptr);
    }
    return result;
}

// Add a transition barrier to cmd_list. Returns true if cmd_list was modified.
bool AddTransitionBarrier(ID3D12GraphicsCommandList*     cmd_list,
                          ID3D12Resource*                resource,
                          UINT                           subresource,
                          const dx12::ResourceStateInfo& current_state,
                          const dx12::ResourceStateInfo& desired_state)
{
    // It is expected that barrier_flags can have bits set for either D3D12_RESOURCE_BARRIER_FLAG_BEGIN_ONLY or
    // D3D12_RESOURCE_BARRIER_FLAG_END_ONLY, not both.
    GFXRECON_ASSERT((current_state.barrier_flags &
                     (D3D12_RESOURCE_BARRIER_FLAG_BEGIN_ONLY | D3D12_RESOURCE_BARRIER_FLAG_END_ONLY)) !=
                    (D3D12_RESOURCE_BARRIER_FLAG_BEGIN_ONLY | D3D12_RESOURCE_BARRIER_FLAG_END_ONLY));
    GFXRECON_ASSERT((current_state.barrier_flags &
                     (D3D12_RESOURCE_BARRIER_FLAG_BEGIN_ONLY | D3D12_RESOURCE_BARRIER_FLAG_END_ONLY)) !=
                    (D3D12_RESOURCE_BARRIER_FLAG_BEGIN_ONLY | D3D12_RESOURCE_BARRIER_FLAG_END_ONLY));

    // Return early if the resource is already in the desired state
    bool current_barrier_begins_split = ((current_state.barrier_flags & D3D12_RESOURCE_BARRIER_FLAG_BEGIN_ONLY) ==
                                         D3D12_RESOURCE_BARRIER_FLAG_BEGIN_ONLY);
    bool desired_barrier_begins_split = ((desired_state.barrier_flags & D3D12_RESOURCE_BARRIER_FLAG_BEGIN_ONLY) ==
                                         D3D12_RESOURCE_BARRIER_FLAG_BEGIN_ONLY);
    if (current_state.states == desired_state.states)
    {
        if (current_barrier_begins_split == desired_barrier_begins_split)
        {
            return false;
        }
    }

    D3D12_RESOURCE_BARRIER_FLAGS flags =
        desired_barrier_begins_split ? D3D12_RESOURCE_BARRIER_FLAG_BEGIN_ONLY : D3D12_RESOURCE_BARRIER_FLAG_NONE;

    // Handle the cases where the resource is currently in the middle of a split transition barrier.
    if (current_barrier_begins_split)
    {
        // If the previous state started a split transition to the desired state, complete the current split transition.
        if (desired_state.states == current_state.states)
        {
            flags = D3D12_RESOURCE_BARRIER_FLAG_END_ONLY;
        }
        else
        {
            // If the previous state started a split transition to a state other than the desired state, add an
            // additional transition to complete the split transition.
            AddTransitionBarrier(cmd_list,
                                 resource,
                                 subresource,
                                 current_state,
                                 { current_state.states, D3D12_RESOURCE_BARRIER_FLAG_END_ONLY });
        }
    }

    D3D12_RESOURCE_TRANSITION_BARRIER transition;
    transition.pResource   = resource;
    transition.Subresource = subresource;
    transition.StateBefore = current_state.states;
    transition.StateAfter  = desired_state.states;

    D3D12_RESOURCE_BARRIER barrier;
    barrier.Type       = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
    barrier.Flags      = flags;
    barrier.Transition = transition;

    cmd_list->ResourceBarrier(1, &barrier);

    return true;
}

size_t Dx12ResourceDataUtil::GetSubresourceCount(ID3D12Resource* resource)
{
    GFXRECON_ASSERT(resource != nullptr);

    size_t sub_count = 0;
    auto   device    = dx12::GetDeviceComPtrFromChild<ID3D12Device>(resource);
    auto   res_desc  = resource->GetDesc();

    if (res_desc.Dimension == D3D12_RESOURCE_DIMENSION_BUFFER)
    {
        sub_count = 1;
    }
    else
    {
        GFXRECON_ASSERT((res_desc.Dimension == D3D12_RESOURCE_DIMENSION_TEXTURE1D) ||
                        (res_desc.Dimension == D3D12_RESOURCE_DIMENSION_TEXTURE2D) ||
                        (res_desc.Dimension == D3D12_RESOURCE_DIMENSION_TEXTURE3D));

        uint32_t plane_count = 1;

        // Get the plane count for the texture format.  With D3D12, each plane has its own subresource.
        D3D12_FEATURE_DATA_FORMAT_INFO format_info = { res_desc.Format, 0 };
        if (SUCCEEDED(device->CheckFeatureSupport(D3D12_FEATURE_FORMAT_INFO, &format_info, sizeof(format_info))))
        {
            plane_count = format_info.PlaneCount;
        }

        sub_count = res_desc.MipLevels * plane_count;

        if ((res_desc.Dimension == D3D12_RESOURCE_DIMENSION_TEXTURE2D) ||
            (res_desc.Dimension == D3D12_RESOURCE_DIMENSION_TEXTURE1D))
        {
            sub_count *= res_desc.DepthOrArraySize;
        }
    }
    return sub_count;
}

void Dx12ResourceDataUtil::GetResourceCopyInfo(ID3D12Resource*                                  resource,
                                               size_t&                                          subresource_count,
                                               std::vector<uint64_t>&                           subresource_offsets,
                                               std::vector<uint64_t>&                           subresource_sizes,
                                               std::vector<D3D12_PLACED_SUBRESOURCE_FOOTPRINT>& layouts,
                                               uint64_t&                                        total_size,
                                               ID3D12Device*                                    device)
{
    GFXRECON_ASSERT(resource != nullptr);

    dx12::ID3D12DeviceComPtr device_com_ptr;
    if (device == nullptr)
    {
        device_com_ptr = dx12::GetDeviceComPtrFromChild<ID3D12Device>(resource);
        device         = device_com_ptr.GetInterfacePtr();
    }
    GFXRECON_ASSERT(device != nullptr);

    subresource_count = GetSubresourceCount(resource);
    subresource_offsets.clear();
    subresource_sizes.clear();
    layouts.clear();
    total_size = 0;

    auto resource_desc = resource->GetDesc();

    if (resource_desc.Dimension == D3D12_RESOURCE_DIMENSION_BUFFER)
    {
        layouts.resize(subresource_count);
        subresource_sizes.push_back(resource_desc.Width);
        subresource_offsets.push_back(0);

        device->GetCopyableFootprints(&resource_desc, 0, 1, 0, layouts.data(), nullptr, nullptr, &total_size);

        // Total resource size should be equal to buffer width and should have 0 offset.
        GFXRECON_ASSERT((total_size == resource_desc.Width) && (layouts[0].Offset == 0));
    }
    else
    {
        layouts.resize(subresource_count);
        std::vector<UINT>   subresource_row_counts(subresource_count);
        std::vector<UINT64> subresource_row_size_bytes(subresource_count);

        graphics::dx12::RobustGetCopyableFootprint(device,
                                                   resource,
                                                   &resource_desc,
                                                   0,
                                                   subresource_count,
                                                   0,
                                                   layouts.data(),
                                                   subresource_row_counts.data(),
                                                   subresource_row_size_bytes.data(),
                                                   &total_size);

        subresource_sizes.resize(subresource_count);
        subresource_offsets.resize(subresource_count);
        for (size_t i = 0; i < subresource_count; ++i)
        {
            size_t slice_count = 1;
            if (resource_desc.Dimension == D3D12_RESOURCE_DIMENSION_TEXTURE3D)
            {
                slice_count = layouts[i].Footprint.Depth;
            }
            else
            {
                GFXRECON_ASSERT(layouts[i].Footprint.Depth == 1);
            }

            // Compute the size required to copy the subresource. The final row of the subresource does not require
            // padding, so the total size is (row_count - 1) * row_pitch + row_size_bytes.
            size_t row_count     = subresource_row_counts[i] * slice_count;
            subresource_sizes[i] = (row_count - 1) * layouts[i].Footprint.RowPitch + subresource_row_size_bytes[i];

            subresource_offsets[i] = layouts[i].Offset;
        }

        GFXRECON_ASSERT((subresource_offsets.back() + subresource_sizes.back()) == total_size);
    }
}

Dx12ResourceDataUtil::Dx12ResourceDataUtil(ID3D12Device* device, uint64_t min_buffer_size) :
    device_(device), staging_buffers_{ nullptr, nullptr }, staging_buffer_sizes_{ 0, 0 },
    min_buffer_size_(min_buffer_size), fence_value_(0)
{
    HRESULT result = E_FAIL;

    // Create a command list (and its dependencies) that will be used to run the copy commands.
    D3D12_COMMAND_LIST_TYPE  list_type  = D3D12_COMMAND_LIST_TYPE_DIRECT;
    D3D12_COMMAND_QUEUE_DESC queue_desc = {};
    queue_desc.Flags                    = D3D12_COMMAND_QUEUE_FLAG_NONE;
    queue_desc.Type                     = list_type;
    result                              = device_->CreateCommandQueue(&queue_desc, IID_PPV_ARGS(&command_queue_));
    if (SUCCEEDED(result))
    {
        result = device_->CreateCommandAllocator(list_type, IID_PPV_ARGS(&command_allocator_));
        if (SUCCEEDED(result))
        {
            result =
                device_->CreateCommandList(0, list_type, command_allocator_, nullptr, IID_PPV_ARGS(&command_list_));
            if (SUCCEEDED(result))
            {
                result = command_list_->Close();
                if (SUCCEEDED(result))
                {
                    result = device_->CreateFence(fence_value_, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&command_fence_));
                }
            }
        }
    }

    if (FAILED(result))
    {
        GFXRECON_LOG_ERROR(
            "Failed to initialize DX12 command list and command queue for Dx12ResourceDataUtil. (error = %lx)", result);
    }
}

dx12::ID3D12ResourceComPtr Dx12ResourceDataUtil::GetStagingBuffer(CopyType type, uint64_t required_buffer_size)
{
    if (required_buffer_size == 0)
    {
        return E_INVALIDARG;
    }

    // If the buffer was already initialized with a size >= required_buffer_size, use the existing buffer.
    if ((staging_buffers_[type] != nullptr) && (required_buffer_size <= staging_buffer_sizes_[type]))
    {
        return staging_buffers_[type];
    }

    uint64_t buffer_size = std::max(required_buffer_size, min_buffer_size_);

    // Release existing buffer.
    staging_buffers_[type]      = nullptr;
    staging_buffer_sizes_[type] = 0;

    // Create the staging resource.
    D3D12_HEAP_PROPERTIES heap_properties = {};
    if (type == kCopyTypeRead)
    {
        heap_properties.Type = D3D12_HEAP_TYPE_READBACK;
    }
    else
    {
        heap_properties.Type = D3D12_HEAP_TYPE_UPLOAD;
    }

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

    D3D12_RESOURCE_STATES staging_resource_state;
    if (type == kCopyTypeRead)
    {
        staging_resource_state = D3D12_RESOURCE_STATE_COPY_DEST;
    }
    else
    {
        staging_resource_state = D3D12_RESOURCE_STATE_GENERIC_READ;
    }
    HRESULT result = device_->CreateCommittedResource(&heap_properties,
                                                      D3D12_HEAP_FLAG_NONE,
                                                      &staging_resource_desc,
                                                      staging_resource_state,
                                                      nullptr,
                                                      IID_PPV_ARGS(&staging_buffers_[type]));

    if (FAILED(result))
    {
        GFXRECON_LOG_FATAL("Failed to create a staging buffer of size %" PRIu64 " for copying resource data.",
                           buffer_size);

        staging_buffers_[type]      = nullptr;
        staging_buffer_sizes_[type] = 0;
    }
    else
    {
        staging_buffer_sizes_[type] = buffer_size;
    }

    return staging_buffers_[type];
}

dx12::ID3D12ResourceComPtr
Dx12ResourceDataUtil::CreateStagingBuffer(ID3D12Device* device, CopyType type, uint64_t required_buffer_size)
{
    if (required_buffer_size == 0)
    {
        return nullptr;
    }

    // Create the staging resource.
    D3D12_HEAP_TYPE       staging_resource_type  = D3D12_HEAP_TYPE_READBACK;
    D3D12_RESOURCE_STATES staging_resource_state = D3D12_RESOURCE_STATE_COPY_DEST;
    if (type == kCopyTypeWrite)
    {
        staging_resource_type  = D3D12_HEAP_TYPE_UPLOAD;
        staging_resource_state = D3D12_RESOURCE_STATE_GENERIC_READ;
    }
    return graphics::dx12::CreateBufferResource(
        device, required_buffer_size, staging_resource_type, staging_resource_state, D3D12_RESOURCE_FLAG_NONE);
}

dx12::ID3D12ResourceComPtr Dx12ResourceDataUtil::CreateStagingBuffer(CopyType type, uint64_t required_buffer_size)
{
    return CreateStagingBuffer(device_, type, required_buffer_size);
}

HRESULT Dx12ResourceDataUtil::ReadFromResource(ID3D12Resource*                             target_resource,
                                               bool                                        try_map_and_copy,
                                               const std::vector<dx12::ResourceStateInfo>& before_states,
                                               const std::vector<dx12::ResourceStateInfo>& after_states,
                                               std::vector<uint8_t>&                       data,
                                               std::vector<uint64_t>&                      subresource_offsets,
                                               std::vector<uint64_t>&                      subresource_sizes,
                                               ID3D12Resource*                             staging_buffer_for_batching,
                                               ID3D12CommandQueue*                         queue)
{
    HRESULT result = E_FAIL;

    // Get the layout information for copying resource data.
    size_t   subresource_count;
    uint64_t required_data_size;
    temp_subresource_layouts_.clear();
    GetResourceCopyInfo(target_resource,
                        subresource_count,
                        subresource_offsets,
                        subresource_sizes,
                        temp_subresource_layouts_,
                        required_data_size,
                        device_);

    GFXRECON_CHECK_CONVERSION_DATA_LOSS(size_t, required_data_size);

    data.clear();

    // If the resource can be mapped, map it, copy the data, and return success.
    if (try_map_and_copy && IsResourceCpuAccessible(target_resource, kCopyTypeRead))
    {
        data.resize(static_cast<size_t>(required_data_size), 0);
        if (CopyMappableResource(
                target_resource, kCopyTypeRead, &data, nullptr, subresource_offsets, subresource_sizes))
        {
            return S_OK;
        }
    }

    // Get or create the staging buffer for copying the resource on device.
    bool            batching;
    ID3D12Resource* staging_resource;
    if (staging_buffer_for_batching != nullptr)
    {
        batching         = true;
        staging_resource = staging_buffer_for_batching;
    }
    else
    {
        batching         = false;
        staging_resource = GetStagingBuffer(kCopyTypeRead, required_data_size);
    }

    // Build and execute commands to copy data to the resource.
    result = ExecuteCopyCommandList(target_resource,
                                    kCopyTypeRead,
                                    required_data_size,
                                    temp_subresource_layouts_,
                                    before_states,
                                    after_states,
                                    staging_resource,
                                    batching,
                                    queue);

    // After the command list has completed, map the copy resource and read its data.
    if (!batching && SUCCEEDED(result))
    {
        data.resize(static_cast<size_t>(required_data_size), 0);
        result = MapSubresourceAndReadData(staging_resource, 0, static_cast<size_t>(required_data_size), data.data());
    }

    return result;
}

HRESULT Dx12ResourceDataUtil::WriteToResource(ID3D12Resource*                             target_resource,
                                              bool                                        try_map_and_copy,
                                              const std::vector<dx12::ResourceStateInfo>& before_states,
                                              const std::vector<dx12::ResourceStateInfo>& after_states,
                                              const std::vector<uint8_t>&                 data,
                                              const std::vector<uint64_t>&                subresource_offsets,
                                              const std::vector<uint64_t>&                subresource_sizes,
                                              ID3D12Resource*                             staging_buffer_for_batching)
{
    HRESULT result = E_FAIL;

    size_t   subresource_count;
    uint64_t required_data_size;

    // Get the layout information for copying resource data.
    std::vector<uint64_t> layout_offsets;
    std::vector<uint64_t> layout_sizes;
    temp_subresource_layouts_.clear();
    GetResourceCopyInfo(target_resource,
                        subresource_count,
                        layout_offsets,
                        layout_sizes,
                        temp_subresource_layouts_,
                        required_data_size,
                        device_);

    // If the resource can be mapped, map it, copy the data, and return success.
    if (try_map_and_copy && IsResourceCpuAccessible(target_resource, kCopyTypeWrite))
    {
        if (CopyMappableResource(
                target_resource, kCopyTypeWrite, nullptr, &data, subresource_offsets, subresource_sizes))
        {
            return S_OK;
        }
    }

    // Get or create the staging buffer for copying the resource on device.
    bool            batching;
    ID3D12Resource* staging_resource;
    if (staging_buffer_for_batching != nullptr)
    {
        batching         = true;
        staging_resource = staging_buffer_for_batching;
    }
    else
    {
        batching         = false;
        staging_resource = GetStagingBuffer(kCopyTypeWrite, required_data_size);
    }

    // When writing and before running the command list, map the copy resource and write its data.
    uint8_t* subresource_data;
    result = dx12::MapSubresource(staging_resource, 0, &dx12::kZeroRange, subresource_data);
    if (SUCCEEDED(result))
    {
        for (UINT i = 0; i < subresource_count; ++i)
        {
            GFXRECON_CHECK_CONVERSION_DATA_LOSS(size_t, subresource_sizes[i]);
            size_t subresource_size = static_cast<size_t>(std::min(subresource_sizes[i], layout_sizes[i]));

            if (layout_sizes[i] > subresource_sizes[i])
            {
                GFXRECON_LOG_DEBUG("The size of the data to be copied to the subresource does not match the size "
                                   "required by the subresource's copyable footprint (data size = %" PRIu64
                                   ", footprint size = %" PRIu64 ", subresouce index = %" PRIu32 ").",
                                   subresource_sizes[i],
                                   layout_sizes[i],
                                   i);
            }
            else if (layout_sizes[i] < subresource_sizes[i])
            {
                GFXRECON_LOG_ERROR("The size of the data to be copied to the subresource is greater than the size of "
                                   "the subresource's copyable footprint."
                                   "Not all data can be written to the subresource. (data size = %" PRIu64
                                   ", footprint size = %" PRIu64 ", subresouce index = %" PRIu32 ").",
                                   subresource_sizes[i],
                                   layout_sizes[i],
                                   i);
            }
            size_t layout_offset = static_cast<size_t>(layout_offsets[i]);
            size_t layout_size   = static_cast<size_t>(layout_sizes[i]);
            util::platform::MemoryCopy(
                subresource_data + layout_offset, layout_size, data.data() + subresource_offsets[i], subresource_size);
        }
        staging_resource->Unmap(0, nullptr);
    }
    else
    {
        GFXRECON_LOG_ERROR("Failed to map staging buffer for writing data to resource.");
    }

    // Build and execute commands to copy data to the resource.
    result = ExecuteCopyCommandList(target_resource,
                                    kCopyTypeWrite,
                                    required_data_size,
                                    temp_subresource_layouts_,
                                    before_states,
                                    after_states,
                                    staging_resource,
                                    batching);

    return result;
}

bool Dx12ResourceDataUtil::IsResourceCpuAccessible(ID3D12Resource* resource, CopyType copy_type)
{
    bool                  cpu_accessible = false;
    D3D12_HEAP_PROPERTIES heap_properties;
    D3D12_HEAP_FLAGS      heap_flags;
    HRESULT               hr = resource->GetHeapProperties(&heap_properties, &heap_flags);
    if (SUCCEEDED(hr))
    {
        cpu_accessible = ((heap_properties.Type == D3D12_HEAP_TYPE_UPLOAD) && (copy_type == kCopyTypeWrite)) ||
                         ((heap_properties.Type == D3D12_HEAP_TYPE_READBACK) && (copy_type == kCopyTypeRead)) ||
                         ((heap_properties.Type == D3D12_HEAP_TYPE_CUSTOM) &&
                          (heap_properties.CPUPageProperty != D3D12_CPU_PAGE_PROPERTY_UNKNOWN) &&
                          (heap_properties.CPUPageProperty != D3D12_CPU_PAGE_PROPERTY_NOT_AVAILABLE));
    }
    else
    {
        GFXRECON_LOG_DEBUG("Call to GetHeapProperties failed. Unable to determine if resource is CPU accessible.");
    }
    return cpu_accessible;
}

bool Dx12ResourceDataUtil::CopyMappableResource(ID3D12Resource*              target_resource,
                                                CopyType                     copy_type,
                                                std::vector<uint8_t>*        read_data,
                                                const std::vector<uint8_t>*  write_data,
                                                const std::vector<uint64_t>& subresource_offsets,
                                                const std::vector<uint64_t>& subresource_sizes)
{
    uint64_t subresource_count = subresource_offsets.size();

    GFXRECON_ASSERT((copy_type != kCopyTypeRead) || (read_data != nullptr));
    GFXRECON_ASSERT((copy_type != kCopyTypeWrite) || (write_data != nullptr));

    HRESULT result = S_OK;
    for (UINT i = 0; i < subresource_count && SUCCEEDED(result); ++i)
    {
        GFXRECON_CHECK_CONVERSION_DATA_LOSS(size_t, subresource_sizes[i]);
        size_t subresource_size = static_cast<size_t>(subresource_sizes[i]);

        if (copy_type == kCopyTypeRead)
        {
            result = MapSubresourceAndReadData(
                target_resource, i, subresource_size, read_data->data() + subresource_offsets[i]);
        }
        else
        {
            result = MapSubresourceAndWriteData(
                target_resource, i, subresource_size, write_data->data() + subresource_offsets[i]);
        }
    }

    return SUCCEEDED(result);
}

HRESULT Dx12ResourceDataUtil::ExecuteAndWaitForCommandList(ID3D12CommandQueue* queue)
{
    // Execute the command list and wait for completion.
    ID3D12CommandList* cmd_lists[] = { command_list_ };

    ID3D12CommandQueue* command_queue = queue;
    if (command_queue == nullptr)
    {
        command_queue = command_queue_;
    }
    command_queue->ExecuteCommandLists(1, cmd_lists);
    return dx12::WaitForQueue(command_queue, command_fence_, ++fence_value_);

    // MakeResident and Evict are ref-counted. Remove the ref count added by MakeResident.
    for (auto resource : resident_resources)
    {
        if (!SUCCEEDED(device_->Evict(1, &resource)))
        {
            GFXRECON_LOG_WARNING("Failed to evict resource after copying resource data.");
        }
    }
}

HRESULT Dx12ResourceDataUtil::CloseCommandList()
{
    return command_list_->Close();
}

HRESULT
Dx12ResourceDataUtil::ResetCommandList()
{
    HRESULT result = command_allocator_->Reset();
    if (SUCCEEDED(result))
    {
        result = command_list_->Reset(command_allocator_, nullptr);
        if (FAILED(result))
        {
            GFXRECON_LOG_ERROR("Could not reset command list to copy resource data. (error = %lx)", result);
        }
    }
    else
    {
        GFXRECON_LOG_ERROR("Could not reset command allocator to copy resource data.. (error = %lx)", result);
    }
    return result;
}

HRESULT
Dx12ResourceDataUtil::RecordCommandsToCopyResource(
    ID3D12GraphicsCommandList*                             target_command_list,
    ID3D12Resource*                                        target_resource,
    CopyType                                               copy_type,
    uint64_t                                               copy_size,
    const std::vector<D3D12_PLACED_SUBRESOURCE_FOOTPRINT>& subresource_layouts,
    const std::vector<dx12::ResourceStateInfo>&            before_states,
    const std::vector<dx12::ResourceStateInfo>&            after_states,
    ID3D12Resource*                                        staging_buffer)
{
    const dx12::ResourceStateInfo common_state = { D3D12_RESOURCE_STATE_COMMON, D3D12_RESOURCE_BARRIER_FLAG_NONE };
    const dx12::ResourceStateInfo copy_state   = { copy_type == kCopyTypeRead ? D3D12_RESOURCE_STATE_COPY_SOURCE
                                                                              : D3D12_RESOURCE_STATE_COPY_DEST,
                                                 D3D12_RESOURCE_BARRIER_FLAG_NONE };

    uint64_t subresource_count = subresource_layouts.size();

    // The number of incoming subresource states should match the number of subresource layouts.
    if ((subresource_count != before_states.size()) || (subresource_count != after_states.size()))
    {
        GFXRECON_LOG_ERROR("Unexpected number of subresource states for copying resource data.");
        return E_INVALIDARG;
    }

    // Transition the resource, copy the data, and transition it back.
    auto resource_desc = target_resource->GetDesc();

    if (resource_desc.SampleDesc.Count > 1)
    {
        GFXRECON_LOG_ERROR("Dx12ResourceDataUtil: Multi-sampled resources are not supported.");
        return E_INVALIDARG;
    }

    // Transition subresources to copy state.
    for (UINT i = 0; i < subresource_count; ++i)
    {
        if ((before_states[i].states & D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE) ==
            D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE)
        {
            GFXRECON_LOG_ERROR("Copying ray tracing acceleration structure resources is not supported. Ray tracing "
                               "acceleration structure resources cannot be transitioned out of resource state "
                               "D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE.");
            return E_INVALIDARG;
        }

        // Prepare resource state. Resources in common state will be implicitly promoted to the appropriate copy state.
        AddTransitionBarrier(target_command_list, target_resource, i, before_states[i], common_state);
    }

    // Copy subresources' data.
    for (UINT i = 0; i < subresource_count; ++i)
    {
        // Copy data.
        if (resource_desc.Dimension == D3D12_RESOURCE_DIMENSION_BUFFER)
        {
            // There is only 1 subresource for buffers.
            GFXRECON_ASSERT(i == 0);

            if (copy_type == kCopyTypeRead)
            {
                // Copy from buffer.
                target_command_list->CopyBufferRegion(staging_buffer, 0, target_resource, 0, copy_size);
            }
            else
            {
                // Copy to buffer.
                target_command_list->CopyBufferRegion(target_resource, 0, staging_buffer, 0, copy_size);
            }
        }
        else
        {
            D3D12_TEXTURE_COPY_LOCATION texture_location = {};
            texture_location.pResource                   = target_resource;
            texture_location.Type                        = D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX;
            texture_location.SubresourceIndex            = i;

            D3D12_TEXTURE_COPY_LOCATION copy_location = {};
            copy_location.pResource                   = staging_buffer;
            copy_location.Type                        = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT;
            copy_location.PlacedFootprint             = subresource_layouts[i];

            if (copy_type == kCopyTypeRead)
            {
                // Copy from texture.
                target_command_list->CopyTextureRegion(&copy_location, 0, 0, 0, &texture_location, nullptr);
            }
            else
            {
                // Copy to texture.
                target_command_list->CopyTextureRegion(&texture_location, 0, 0, 0, &copy_location, nullptr);
            }
        }
    }

    // Transition subresources to final state.
    for (UINT i = 0; i < subresource_count; ++i)
    {
        // If the final desired state is STATE_COMMON and the resource will decay to STATE_COMMON, skip the final
        // transition barrier.
        bool decays_to_common = (copy_state.states == D3D12_RESOURCE_STATE_COPY_SOURCE) ||
                                (resource_desc.Dimension == D3D12_RESOURCE_DIMENSION_BUFFER) ||
                                ((resource_desc.Flags & D3D12_RESOURCE_FLAG_ALLOW_SIMULTANEOUS_ACCESS) != 0);
        if (!decays_to_common || (after_states[i].states != D3D12_RESOURCE_STATE_COMMON))
        {
            AddTransitionBarrier(target_command_list, target_resource, i, copy_state, after_states[i]);
        }
    }

    return S_OK;
}

HRESULT
Dx12ResourceDataUtil::ExecuteCopyCommandList(ID3D12Resource*                                        target_resource,
                                             CopyType                                               copy_type,
                                             uint64_t                                               copy_size,
                                             const std::vector<D3D12_PLACED_SUBRESOURCE_FOOTPRINT>& subresource_layouts,
                                             const std::vector<dx12::ResourceStateInfo>&            before_states,
                                             const std::vector<dx12::ResourceStateInfo>&            after_states,
                                             ID3D12Resource*                                        staging_buffer,
                                             bool                                                   batching,
                                             ID3D12CommandQueue*                                    queue)
{
    // Make sure the target resource is resident.
    ID3D12Pageable* const pageable = target_resource;
    if (!SUCCEEDED(device_->MakeResident(1, &pageable)))
    {
        GFXRECON_LOG_ERROR("Failed to make resource resident for copying resource data.");
        return E_FAIL;
    }

    GFXRECON_ASSERT(staging_buffer != nullptr);

    resident_resources.push_back(pageable);
    // Transition the resource, copy the data, and transition it back.
    auto resource_desc = target_resource->GetDesc();

    HRESULT result = S_OK;
    if (batching == false)
    {
        result = command_allocator_->Reset();
    }
    if (SUCCEEDED(result))
    {
        if (batching == false)
        {
            result = command_list_->Reset(command_allocator_, nullptr);
        }
        if (SUCCEEDED(result))
        {
            result = RecordCommandsToCopyResource(command_list_,
                                                  target_resource,
                                                  copy_type,
                                                  copy_size,
                                                  subresource_layouts,
                                                  before_states,
                                                  after_states,
                                                  staging_buffer);

            // Close and execute the command list.
            if (SUCCEEDED(result) && (batching == false))
            {
                result = command_list_->Close();
                if (SUCCEEDED(result))
                {
                    result = ExecuteAndWaitForCommandList(queue);
                }
            }
        }
    }

    if (FAILED(result))
    {
        GFXRECON_LOG_ERROR("Error executing command list to copy resource data. (error = %lx)", result);
    }

    return result;
}

void Dx12ResourceDataUtil::InitializeMetaCommand(ID3D12MetaCommand* pMetaCommand,
                                                 const void*        pInitializationParametersData,
                                                 SIZE_T             InitializationParametersDataSizeInBytes)
{
    command_list_->InitializeMetaCommand(
        pMetaCommand, pInitializationParametersData, InitializationParametersDataSizeInBytes);
}

HRESULT
Dx12ResourceDataUtil::ExecuteTransitionCommandList(ID3D12Resource*                             target_resource,
                                                   const std::vector<dx12::ResourceStateInfo>& before_states,
                                                   const std::vector<dx12::ResourceStateInfo>& after_states,
                                                   ID3D12CommandQueue*                         queue)
{
    GFXRECON_ASSERT(before_states.size() == after_states.size());
    uint64_t subresource_count = before_states.size();

    // Transition the subresources.
    bool    cmd_list_empty = true;
    auto    resource_desc  = target_resource->GetDesc();
    HRESULT result         = command_allocator_->Reset();
    if (SUCCEEDED(result))
    {
        result = command_list_->Reset(command_allocator_, nullptr);
        if (SUCCEEDED(result))
        {
            for (UINT i = 0; i < subresource_count; ++i)
            {
                if (AddTransitionBarrier(command_list_, target_resource, i, before_states[i], after_states[i]))
                {
                    cmd_list_empty = false;
                }
            }

            // Close and execute the command list.
            result = command_list_->Close();
            if (SUCCEEDED(result) && !cmd_list_empty)
            {
                result = ExecuteAndWaitForCommandList(queue);
            }
        }
    }

    if (FAILED(result))
    {
        GFXRECON_LOG_ERROR("Error executing command list to transition resource state. (error = %lx)", result);
    }

    return result;
}

GFXRECON_END_NAMESPACE(graphics)
GFXRECON_END_NAMESPACE(gfxrecon)
