/*
 * Copyright (c) Contributors to the Open 3D Engine Project.
 * For complete copyright and license terms please see the LICENSE at the root of this distribution.
 *
 * SPDX-License-Identifier: Apache-2.0 OR MIT
 *
 */
#include <Atom/RHI.Reflect/ConstantsLayout.h>
#include <AzCore/Serialization/SerializeContext.h>
#include <AzCore/Utils/TypeHash.h>

namespace AZ
{
    namespace RHI
    {
        void ConstantsLayout::Reflect(AZ::ReflectContext* context)
        {
            if (SerializeContext* serializeContext = azrtti_cast<SerializeContext*>(context))
            {
                serializeContext->Class<ConstantsLayout>()
                    ->Version(1)    // Version 1: Adding debug helper functions to Shader Resource Groups
                    ->Field("m_inputs", &ConstantsLayout::m_inputs)
                    ->Field("m_idReflection", &ConstantsLayout::m_idReflection)
                    ->Field("m_sizeInBytes", &ConstantsLayout::m_sizeInBytes)
                    ->Field("m_hash", &ConstantsLayout::m_hash);
            }

            IdReflectionMapForConstants::Reflect(context);
        }

        RHI::Ptr<ConstantsLayout> ConstantsLayout::Create()
        {
            return aznew ConstantsLayout;
        }

        void ConstantsLayout::AddShaderInput(const ShaderInputConstantDescriptor& descriptor)
        {
            m_inputs.push_back(descriptor);
        }

        void ConstantsLayout::Clear()
        {
            m_inputs.clear();
            m_idReflection.Clear();
            m_sizeInBytes = 0;
            m_hash = InvalidHash;
        }

        bool ConstantsLayout::IsFinalized() const
        {
            return m_hash != InvalidHash;
        }

        bool ConstantsLayout::Finalize()
        {
            m_hash = HashValue64{ 0 };
            uint32_t constantInputIndex = 0;
            uint32_t constantDataSize = 0;
            for (const auto& constantDescriptor : m_inputs)
            {
                const ShaderInputConstantIndex inputIndex(constantInputIndex);
                if (!m_idReflection.Insert(constantDescriptor.m_name, inputIndex))
                {
                    Clear();
                    return false;
                }

                uint32_t end = constantDescriptor.m_constantByteOffset + constantDescriptor.m_constantByteCount;
                constantDataSize = AZStd::max(constantDataSize, end);

                ++constantInputIndex;
                m_hash = TypeHash64(constantDescriptor.GetHash(), m_hash);
            }

            m_sizeInBytes = constantDataSize;

            if (!ValidateConstantInputs())
            {
                Clear();
                return false;
            }

            return true;
        }

        HashValue64 ConstantsLayout::GetHash() const
        {
            return m_hash;
        }

        ShaderInputConstantIndex ConstantsLayout::FindShaderInputIndex(const Name& name) const
        {
            return m_idReflection.Find(name);
        }

        Interval ConstantsLayout::GetInterval(ShaderInputConstantIndex inputIndex) const
        {
            const ShaderInputConstantDescriptor& desc = GetShaderInput(inputIndex);
            uint32_t start = desc.m_constantByteOffset;
            uint32_t end = start + desc.m_constantByteCount;
            return Interval(start, end);
        }

        const ShaderInputConstantDescriptor& ConstantsLayout::GetShaderInput(ShaderInputConstantIndex inputIndex) const
        {
            return m_inputs[inputIndex.GetIndex()];
        }

        AZStd::array_view<ShaderInputConstantDescriptor> ConstantsLayout::GetShaderInputList() const
        {
            return m_inputs;
        }

        uint32_t ConstantsLayout::GetDataSize() const
        {
            return m_sizeInBytes;
        }

        bool ConstantsLayout::ValidateAccess(ShaderInputConstantIndex inputIndex) const
        {
            if (Validation::IsEnabled())
            {
                uint32_t count = static_cast<uint32_t>(m_inputs.size());
                if (inputIndex.GetIndex() >= count)
                {
                    AZ_Assert(false, "Inline Constant Input index '%d' out of range [0,%d).", inputIndex.GetIndex(), count);
                    return false;
                }
            }

            return true;
        }

        bool ConstantsLayout::ValidateConstantInputs() const
        {
            if (Validation::IsEnabled())
            {
                if (!m_sizeInBytes)
                {
                    AZ_Assert(m_idReflection.IsEmpty(), "Constants size is not valid.");
                    return m_idReflection.IsEmpty();
                }
            }

            return true;
        }

        void ConstantsLayout::DebugPrintNames(AZStd::array_view<ShaderInputConstantIndex> constantList) const
        {
            AZStd::string output;
            for (const ShaderInputConstantIndex& constantIdx : constantList)
            {
                if (constantIdx.GetIndex() < m_inputs.size())
                {
                    output += m_inputs[constantIdx.GetIndex()].m_name.GetCStr();
                    output += " - ";
                }
            }
            AZ_Printf("RHI", output.c_str());
        }

    }
}