// Copyright 2025 The Dawn & Tint Authors
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its
//    contributors may be used to endorse or promote products derived from
//    this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "dawn/native/vulkan/ResourceTableVk.h"

#include <vector>

#include "dawn/common/Enumerator.h"
#include "dawn/native/DynamicUploader.h"
#include "dawn/native/vulkan/DeviceVk.h"
#include "dawn/native/vulkan/TextureVk.h"
#include "dawn/native/vulkan/UtilsVulkan.h"
#include "dawn/native/vulkan/VulkanError.h"

namespace dawn::native::vulkan {

// static
ResultOrError<Ref<ResourceTable>> ResourceTable::Create(Device* device,
                                                        const ResourceTableDescriptor* descriptor) {
    Ref<ResourceTable> table = AcquireRef(new ResourceTable(device, descriptor));
    DAWN_TRY(table->Initialize());
    return table;
}

// static
ResultOrError<VkDescriptorSetLayout> ResourceTable::MakeDescriptorSetLayout(Device* device) {
    // A resource table is a bindgroup made of two entries:
    //
    //  - Binding 0: the metadata storage buffer.
    //  - Binding 1: the variable length, partially bound, update-after-bind array of sampled
    //  textures/samplers.
    std::array<VkDescriptorSetLayoutBinding, 2> bindings = {
        {{
             .binding = 0,
             .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
             .descriptorCount = 1,
             .stageFlags = VulkanShaderStages(kAllStages),
             .pImmutableSamplers = nullptr,
         },
         {
             .binding = 1,
             .descriptorType = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
             .descriptorCount = device->GetLimits().resourceTableLimits.maxResourceTableSize +
                                uint32_t(GetDefaultResourceCount()),
             .stageFlags = VulkanShaderStages(kAllStages),
             .pImmutableSamplers = nullptr,
         }}};
    std::array<VkDescriptorBindingFlags, 2> flags = {
        0,  //
        VK_DESCRIPTOR_BINDING_UPDATE_AFTER_BIND_BIT |
            VK_DESCRIPTOR_BINDING_UPDATE_UNUSED_WHILE_PENDING_BIT |
            VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT |
            VK_DESCRIPTOR_BINDING_VARIABLE_DESCRIPTOR_COUNT_BIT};

    VkDescriptorSetLayoutBindingFlagsCreateInfo flagCreateInfo{
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_BINDING_FLAGS_CREATE_INFO,
        .pNext = nullptr,
        .bindingCount = flags.size(),
        .pBindingFlags = flags.data(),
    };
    VkDescriptorSetLayoutCreateInfo createInfo{
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        .pNext = &flagCreateInfo,
        .flags = VK_DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT,
        .bindingCount = bindings.size(),
        .pBindings = bindings.data(),
    };

    VkDescriptorSetLayout dsLayout = VK_NULL_HANDLE;
    DAWN_TRY(CheckVkSuccess(device->fn.CreateDescriptorSetLayout(device->GetVkDevice(), &createInfo,
                                                                 nullptr, &*dsLayout),
                            "CreateDescriptorSetLayout"));
    return dsLayout;
}

ResourceTable::~ResourceTable() = default;

MaybeError ResourceTable::Initialize() {
    DAWN_TRY(ResourceTableBase::InitializeBase());

    Device* device = ToBackend(GetDevice());

    uint32_t sampledImageCount = uint32_t(GetSizeWithDefaultResources());

    // Allocate mPool.
    {
        std::array<VkDescriptorPoolSize, 2> sizes = {
            VkDescriptorPoolSize{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1},  // The metadata buffer.
            VkDescriptorPoolSize{VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, sampledImageCount},
        };

        VkDescriptorPoolCreateInfo createInfo{
            .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
            .pNext = nullptr,
            .flags = VK_DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT,
            .maxSets = 1,
            .poolSizeCount = uint32_t(sizes.size()),
            .pPoolSizes = sizes.data(),
        };

        DAWN_TRY(CheckVkOOMThenSuccess(
            device->fn.CreateDescriptorPool(device->GetVkDevice(), &createInfo, nullptr, &*mPool),
            "CreateDescriptorPool"));
    }

    // Allocate mSet from mPool.
    {
        VkDescriptorSetVariableDescriptorCountAllocateInfo variableCountInfo{
            .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_VARIABLE_DESCRIPTOR_COUNT_ALLOCATE_INFO,
            .pNext = nullptr,
            .descriptorSetCount = 1,
            .pDescriptorCounts = &sampledImageCount,
        };
        VkDescriptorSetAllocateInfo allocateInfo{
            .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
            .pNext = &variableCountInfo,
            .descriptorPool = mPool,
            .descriptorSetCount = 1,
            .pSetLayouts = &device->GetResourceTableLayout().GetHandle(),
        };

        DAWN_TRY(CheckVkOOMThenSuccess(
            device->fn.AllocateDescriptorSets(device->GetVkDevice(), &allocateInfo, &*mSet),
            "AllocateDescriptorSets"));
    }

    // Only write the metadata buffer in mSet initially, all the other bindings will be written as
    // needed when they are inserted in the ResourceTable.
    {
        Buffer* metadataBuffer = ToBackend(GetMetadataBuffer());
        VkDescriptorBufferInfo bufferInfo = {
            .buffer = metadataBuffer->GetHandle(),
            .offset = 0,
            .range = metadataBuffer->GetSize(),
        };
        VkWriteDescriptorSet write = {
            .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
            .pNext = nullptr,
            .dstSet = mSet,
            .dstBinding = 0,
            .dstArrayElement = 0,
            .descriptorCount = 1,
            .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            .pImageInfo = nullptr,
            .pBufferInfo = &bufferInfo,
            .pTexelBufferView = nullptr,
        };
        device->fn.UpdateDescriptorSets(device->GetVkDevice(), 1, &write, 0, nullptr);
    }

    return {};
}

// Apply updates to resources or to the metadata buffers that are pending.
MaybeError ResourceTable::ApplyPendingUpdates(CommandRecordingContext* recordingContext) {
    Updates updates = AcquireDirtySlotUpdates();

    if (!updates.metadataUpdates.empty()) {
        DAWN_TRY(UpdateMetadataBuffer(recordingContext, updates.metadataUpdates));
    }
    if (!updates.resourceUpdates.empty()) {
        UpdateResourceBindings(updates.resourceUpdates);
    }

    return {};
}

MaybeError ResourceTable::UpdateMetadataBuffer(CommandRecordingContext* recordingContext,
                                               const std::vector<MetadataUpdate>& updates) {
    // Updates the metadata buffer by scheduling a copy for each u32 that needs to be updated.
    // TODO(https://issues.chromium.org/473442435): If we had a way to use Dawn reentrantly now, we
    // could use a compute shader to dispatch the updates instead of individual copies for each
    // update, and move that logic in the frontend to share it between backends. (also a single
    // dispatch could update multiple metadata buffers potentially).
    Device* device = ToBackend(GetDevice());

    // Allocate enough space for all the data to modify and schedule the copies.
    return device->GetDynamicUploader()->WithUploadReservation(
        sizeof(uint32_t) * updates.size(), kCopyBufferToBufferOffsetAlignment,
        [&](UploadReservation reservation) -> MaybeError {
            uint32_t* stagedData = static_cast<uint32_t*>(reservation.mappedPointer);

            // The metadata buffer will be copied to.
            Buffer* metadataBuffer = ToBackend(GetMetadataBuffer());
            DAWN_ASSERT(metadataBuffer->IsInitialized());
            metadataBuffer->TransitionUsageNow(recordingContext, wgpu::BufferUsage::CopyDst);

            // Prepare the copies.
            std::vector<VkBufferCopy> copies(updates.size());
            for (auto [i, update] : Enumerate(updates)) {
                stagedData[i] = update.data;

                VkBufferCopy copy{
                    .srcOffset = reservation.offsetInBuffer + i * sizeof(uint32_t),
                    .dstOffset = update.offset,
                    .size = sizeof(uint32_t),
                };
                copies[i] = copy;
            }

            // Enqueue the copy commands all at once.
            device->fn.CmdCopyBuffer(recordingContext->commandBuffer,
                                     ToBackend(reservation.buffer)->GetHandle(),
                                     metadataBuffer->GetHandle(), copies.size(), copies.data());

            // Transition the buffer back to be used as storage as that's how it will be used for
            // shader-side validation.
            metadataBuffer->TransitionUsageNow(recordingContext, kReadOnlyStorageBuffer,
                                               kAllStages);

            return {};
        });
}

void ResourceTable::UpdateResourceBindings(const std::vector<ResourceUpdate>& updates) {
    std::vector<VkDescriptorImageInfo> imageWrites;
    std::vector<uint32_t> arrayElements;

    for (const ResourceUpdate& update : updates) {
        // TODO(https://issues.chromium.org/473354063): Support samplers updates.
        // TODO(https://issues.chromium.org/473444515): Support buffer, texel buffers and storage
        // textures.

        VkImageView handle = ToBackend(update.textureView)->GetHandle();
        if (handle == nullptr) {
            continue;
        }

        VkDescriptorImageInfo imageWrite = {
            .sampler = VkSampler{},
            .imageView = handle,
            .imageLayout = VulkanImageLayout(update.textureView->GetFormat(),
                                             wgpu::TextureUsage::TextureBinding),
        };
        imageWrites.push_back(imageWrite);
        arrayElements.push_back(uint32_t(update.slot));
    }

    std::vector<VkWriteDescriptorSet> writes;
    for (size_t i = 0; i < imageWrites.size(); i++) {
        VkWriteDescriptorSet write{
            .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
            .pNext = nullptr,
            .dstSet = mSet,
            .dstBinding = 1,
            .dstArrayElement = arrayElements[i],
            .descriptorCount = 1,
            .descriptorType = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
            .pImageInfo = &imageWrites[i],
            .pBufferInfo = nullptr,
            .pTexelBufferView = nullptr,
        };
        writes.push_back(write);
    }

    Device* device = ToBackend(GetDevice());
    device->fn.UpdateDescriptorSets(device->GetVkDevice(), writes.size(), writes.data(), 0,
                                    nullptr);
}

VkDescriptorSet ResourceTable::GetHandle() const {
    return mSet;
}

void ResourceTable::DestroyImpl(DestroyReason reason) {
    if (mPool != VK_NULL_HANDLE) {
        ToBackend(GetDevice())->GetFencedDeleter()->DeleteWhenUnused(mPool);
        mPool = VK_NULL_HANDLE;
        mSet = VK_NULL_HANDLE;
    }

    ResourceTableBase::DestroyImpl(reason);
}

void ResourceTable::SetLabelImpl() {
    SetDebugName(ToBackend(GetDevice()), mSet, "Dawn_ResourceTable", GetLabel());
}

}  // namespace dawn::native::vulkan