/*
 *  Copyright 2019-2021 Diligent Graphics LLC
 *  Copyright 2015-2019 Egor Yusov
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *  In no event and under no legal theory, whether in tort (including negligence), 
 *  contract, or otherwise, unless required by applicable law (such as deliberate 
 *  and grossly negligent acts) or agreed to in writing, shall any Contributor be
 *  liable for any damages, including any direct, indirect, special, incidental, 
 *  or consequential damages of any character arising as a result of this License or 
 *  out of the use or inability to use the software (including but not limited to damages 
 *  for loss of goodwill, work stoppage, computer failure or malfunction, or any and 
 *  all other commercial damages or losses), even if such Contributor has been advised 
 *  of the possibility of such damages.
 */

#pragma once

/// \file
/// Declaration of Diligent::ShaderResourceCacheVk class

// Shader resource cache stores Vk resources in a continuous chunk of memory:
//
//                                 |Vulkan Descriptor Set|
//                                           A    ___________________________________________________________
//  m_pMemory                                |   |              m_pResources, m_NumResources == m            |
//  |               m_DescriptorSetAllocation|   |                                                           |
//  V                                        |   |                                                           V
//  |  DescriptorSet[0]  |   ....    |  DescriptorSet[Ns-1]  |  Res[0]  |  ... |  Res[n-1]  |    ....     | Res[0]  |  ... |  Res[m-1]  |
//         |    |                                                A \
//         |    |                                                |  \
//         |    |________________________________________________|   \RefCntAutoPtr
//         |             m_pResources, m_NumResources == n            \_________
//         |                                                          |  Object |
//         | m_DescriptorSetAllocation                                 ---------
//         V
//  |Vulkan Descriptor Set|
//
//  Ns = m_NumSets
//
//
// Descriptor set for static and mutable resources is assigned during cache initialization
// Descriptor set for dynamic resources is assigned at every draw call

#include <vector>
#include "DescriptorPoolManager.hpp"
#include "SPIRVShaderResources.hpp"
#include "BufferVkImpl.hpp"
#include "PipelineResourceSignatureVkImpl.hpp"

namespace Diligent
{

class DeviceContextVkImpl;

// sizeof(ShaderResourceCacheVk) == 24 (x64, msvc, Release)
class ShaderResourceCacheVk
{
public:
    using CacheContentType = PipelineResourceSignatureVkImpl::CacheContentType;

    explicit ShaderResourceCacheVk(CacheContentType ContentType) noexcept :
        m_TotalResources{0},
        m_ContentType{static_cast<Uint32>(ContentType)}
    {
        VERIFY_EXPR(GetContentType() == ContentType);
    }

    // clang-format off
    ShaderResourceCacheVk             (const ShaderResourceCacheVk&) = delete;
    ShaderResourceCacheVk             (ShaderResourceCacheVk&&)      = delete;
    ShaderResourceCacheVk& operator = (const ShaderResourceCacheVk&) = delete;
    ShaderResourceCacheVk& operator = (ShaderResourceCacheVk&&)      = delete;
    // clang-format on

    ~ShaderResourceCacheVk();

    static size_t GetRequiredMemorySize(Uint32 NumSets, const Uint32* SetSizes);

    void InitializeSets(IMemoryAllocator& MemAllocator, Uint32 NumSets, const Uint32* SetSizes);
    void InitializeResources(Uint32 Set, Uint32 Offset, Uint32 ArraySize, DescriptorType Type);

    // sizeof(Resource) == 16 (x64, msvc, Release)
    struct Resource
    {
        // clang-format off
        explicit Resource(DescriptorType _Type) noexcept :
            Type{_Type}
        {}

        Resource             (const Resource&) = delete;
        Resource             (Resource&&)      = delete;
        Resource& operator = (const Resource&) = delete;
        Resource& operator = (Resource&&)      = delete;

/* 0 */ const DescriptorType       Type;
/*1-7*/ // Unused
/* 8 */ RefCntAutoPtr<IDeviceObject> pObject;

        VkDescriptorBufferInfo GetUniformBufferDescriptorWriteInfo ()                    const;
        VkDescriptorBufferInfo GetStorageBufferDescriptorWriteInfo ()                    const;
        VkDescriptorImageInfo  GetImageDescriptorWriteInfo  (bool IsImmutableSampler)    const;
        VkBufferView           GetBufferViewWriteInfo       ()                           const;
        VkDescriptorImageInfo  GetSamplerDescriptorWriteInfo()                           const;
        VkDescriptorImageInfo  GetInputAttachmentDescriptorWriteInfo()                   const;
        VkWriteDescriptorSetAccelerationStructureKHR GetAccelerationStructureWriteInfo() const;
        // clang-format on
    };

    // sizeof(DescriptorSet) == 48 (x64, msvc, Release)
    class DescriptorSet
    {
    public:
        // clang-format off
        DescriptorSet(Uint32 NumResources, Resource *pResources) :
            m_NumResources  {NumResources},
            m_pResources    {pResources  }
        {}

        DescriptorSet             (const DescriptorSet&) = delete;
        DescriptorSet             (DescriptorSet&&)      = delete;
        DescriptorSet& operator = (const DescriptorSet&) = delete;
        DescriptorSet& operator = (DescriptorSet&&)      = delete;
        // clang-format on

        inline Resource& GetResource(Uint32 CacheOffset)
        {
            VERIFY(CacheOffset < m_NumResources, "Offset ", CacheOffset, " is out of range");
            return m_pResources[CacheOffset];
        }
        inline const Resource& GetResource(Uint32 CacheOffset) const
        {
            VERIFY(CacheOffset < m_NumResources, "Offset ", CacheOffset, " is out of range");
            return m_pResources[CacheOffset];
        }

        inline Uint32 GetSize() const { return m_NumResources; }

        VkDescriptorSet GetVkDescriptorSet() const
        {
            return m_DescriptorSetAllocation.GetVkDescriptorSet();
        }

        void AssignDescriptorSetAllocation(DescriptorSetAllocation&& Allocation)
        {
            VERIFY(m_NumResources > 0, "Descriptor set is empty");
            m_DescriptorSetAllocation = std::move(Allocation);
        }

        // clang-format off
/* 0 */ const Uint32 m_NumResources = 0;

    private:
/* 8 */ Resource* const m_pResources = nullptr;
/*16 */ DescriptorSetAllocation m_DescriptorSetAllocation;
/*48 */ // End of structure
        // clang-format on
    };

    inline DescriptorSet& GetDescriptorSet(Uint32 Index)
    {
        VERIFY_EXPR(Index < m_NumSets);
        return reinterpret_cast<DescriptorSet*>(m_pMemory)[Index];
    }
    inline const DescriptorSet& GetDescriptorSet(Uint32 Index) const
    {
        VERIFY_EXPR(Index < m_NumSets);
        return reinterpret_cast<const DescriptorSet*>(m_pMemory)[Index];
    }

    inline Uint32 GetNumDescriptorSets() const { return m_NumSets; }
    inline Uint32 GetNumDynamicBuffers() const { return m_NumDynamicBuffers; }

    Uint16& GetDynamicBuffersCounter() { return m_NumDynamicBuffers; }

    CacheContentType GetContentType() const { return static_cast<CacheContentType>(m_ContentType); }

#ifdef DILIGENT_DEBUG
    // Only for debug purposes: indicates what types of resources are stored in the cache
    void DbgVerifyResourceInitialization() const;
    void DbgVerifyDynamicBuffersCounter() const;
#endif

    template <bool VerifyOnly>
    void TransitionResources(DeviceContextVkImpl* pCtxVkImpl);

    __forceinline Uint32 GetDynamicBufferOffsets(Uint32 CtxId, DeviceContextVkImpl* pCtxVkImpl, std::vector<uint32_t>& Offsets) const;

private:
    Resource* GetFirstResourcePtr()
    {
        return reinterpret_cast<Resource*>(reinterpret_cast<DescriptorSet*>(m_pMemory) + m_NumSets);
    }
    const Resource* GetFirstResourcePtr() const
    {
        return reinterpret_cast<const Resource*>(reinterpret_cast<const DescriptorSet*>(m_pMemory) + m_NumSets);
    }

    IMemoryAllocator* m_pAllocator = nullptr;
    void*             m_pMemory    = nullptr;
    Uint16            m_NumSets    = 0;

    // Total actual number of dynamic buffers (that were created with USAGE_DYNAMIC) bound in the resource cache
    // regardless of the variable type. Note this variable is not equal to dynamic offsets count, which is constant.
    Uint16 m_NumDynamicBuffers = 0;
    Uint32 m_TotalResources : 31;

    // Indicates what types of resources are stored in the cache
    const Uint32 m_ContentType : 1;

#ifdef DILIGENT_DEBUG
    // Debug array that stores flags indicating if resources in the cache have been initialized
    std::vector<std::vector<bool>> m_DbgInitializedResources;
#endif
};

__forceinline Uint32 ShaderResourceCacheVk::GetDynamicBufferOffsets(Uint32                 CtxId,
                                                                    DeviceContextVkImpl*   pCtxVkImpl,
                                                                    std::vector<uint32_t>& Offsets) const
{
    // If any of the sets being bound include dynamic uniform or storage buffers, then
    // pDynamicOffsets includes one element for each array element in each dynamic descriptor
    // type binding in each set. Values are taken from pDynamicOffsets in an order such that
    // all entries for set N come before set N+1; within a set, entries are ordered by the binding
    // numbers (unclear if this is SPIRV binding or VkDescriptorSetLayoutBinding number) in the
    // descriptor set layouts; and within a binding array, elements are in order. (13.2.5)

    // In each descriptor set, all uniform buffers with dynamic offsets (DescriptorType::UniformBufferDynamic)
    // for every shader stage come first, followed by all storage buffers with dynamic offsets
    // (DescriptorType::StorageBufferDynamic and DescriptorType::StorageBufferDynamic_ReadOnly) for every shader stage,
    // followed by all other resources.
    Uint32 OffsetInd = 0;
    for (Uint32 set = 0; set < m_NumSets; ++set)
    {
        const auto& DescrSet = GetDescriptorSet(set);
        const auto  SetSize  = DescrSet.GetSize();

        Uint32 res = 0;
        while (res < SetSize)
        {
            const auto& Res = DescrSet.GetResource(res);
            if (Res.Type == DescriptorType::UniformBufferDynamic)
            {
                const auto* pBufferVk = Res.pObject.RawPtr<const BufferVkImpl>();
                auto        Offset    = pBufferVk != nullptr ? pBufferVk->GetDynamicOffset(CtxId, pCtxVkImpl) : 0;
                Offsets[OffsetInd++]  = Offset;
                ++res;
            }
            else
                break;
        }

        while (res < SetSize)
        {
            const auto& Res = DescrSet.GetResource(res);
            if (Res.Type == DescriptorType::StorageBufferDynamic ||
                Res.Type == DescriptorType::StorageBufferDynamic_ReadOnly)
            {
                const auto* pBufferVkView = Res.pObject.RawPtr<const BufferViewVkImpl>();
                const auto* pBufferVk     = pBufferVkView != nullptr ? pBufferVkView->GetBufferVk() : 0;
                auto        Offset        = pBufferVk != nullptr ? pBufferVk->GetDynamicOffset(CtxId, pCtxVkImpl) : 0;
                Offsets[OffsetInd++]      = Offset;
                ++res;
            }
            else
                break;
        }

#ifdef DILIGENT_DEBUG
        for (; res < SetSize; ++res)
        {
            const auto& Res = DescrSet.GetResource(res);
            // clang-format off
            VERIFY(Res.Type != DescriptorType::UniformBufferDynamic && 
                   Res.Type != DescriptorType::StorageBufferDynamic &&
                   Res.Type != DescriptorType::StorageBufferDynamic_ReadOnly, 
                   "All uniform and storage buffers are expected to go first in the beginning of each descriptor set");
            // clang-format on
        }
#endif
    }
    return OffsetInd;
}

} // namespace Diligent
