/*
* All or portions of this file Copyright (c) Amazon.com, Inc. or its affiliates or
* its licensors.
*
* For complete copyright and license terms please see the LICENSE at the root of this
* distribution (the "License"). All use of this software is governed by the License,
* or, if provided, by the license below or the license accompanying this file. Do not
* remove or modify any license notices. This file is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
*
*/

#include <Atom/RPI.Public/Pass/PassUtils.h>
#include <Atom/RPI.Public/Pass/PassFactory.h>
#include <Atom/RPI.Public/Pass/FullscreenTrianglePass.h>
#include <Atom/RPI.Public/RenderPipeline.h>
#include <Atom/RPI.Public/RPIUtils.h>
#include <Atom/RPI.Public/View.h>
#include <Atom/RPI.Public/Shader/ShaderVariant.h>

#include <Atom/RPI.Reflect/Pass/PassTemplate.h>
#include <Atom/RPI.Reflect/Shader/ShaderAsset.h>

#include <Atom/RHI/Factory.h>
#include <Atom/RHI/FrameScheduler.h>
#include <Atom/RHI/PipelineState.h>

#include <AzCore/Asset/AssetCommon.h>
#include <AzCore/Asset/AssetManagerBus.h>

#include <PostProcessing/SMAABasePass.h>

namespace AZ
{
    namespace Render
    {
        SMAABasePass::SMAABasePass(const RPI::PassDescriptor& descriptor)
            : AZ::RPI::FullscreenTrianglePass(descriptor)
        {
        }

        SMAABasePass::~SMAABasePass()
        {
        }

        void SMAABasePass::Init()
        {
            FullscreenTrianglePass::Init();

            AZ_Assert(m_shaderResourceGroup != nullptr, "SMAABasePass %s has a null shader resource group when calling Init.", GetPathName().GetCStr());

            UpdateCurrentShaderVariant();
        }

        void SMAABasePass::CompileResources(const RHI::FrameGraphCompileContext& context)
        {
            AZ_Assert(m_shaderResourceGroup != nullptr, "SMAABasePass %s has a null shader resource group when calling Compile.", GetPathName().GetCStr());

            BindPassSrg(context, m_shaderResourceGroup);

            AZ::Vector4 renderTargetMetrics = CalculateRenderTargetMetrics(GetOutputBinding(0).m_attachment.get());
            if (renderTargetMetrics.GetX() != m_renderTargetMetrics.GetX() ||
                renderTargetMetrics.GetY() != m_renderTargetMetrics.GetY())
            {
                m_renderTargetMetrics = renderTargetMetrics;
                InvalidateSRG();
            }

            if (m_needToUpdateShaderVariant)
            {
                UpdateCurrentShaderVariant();
            }

            if (m_needToUpdateSRG)
            {
                UpdateSRG();

                if (m_shaderResourceGroup->HasShaderVariantKeyFallbackEntry())
                {
                    m_shaderResourceGroup->SetShaderVariantKeyFallbackValue(m_currentShaderVariantKeyFallbackValue);
                }
                m_needToUpdateSRG = false;
            }

            m_shaderResourceGroup->Compile();
        }

        void SMAABasePass::UpdateCurrentShaderVariant()
        {
            auto shaderOption = m_shader->CreateShaderOptionGroup();

            GetCurrentShaderOption(shaderOption);

            m_currentShaderVariantKeyFallbackValue = shaderOption.GetShaderVariantKeyFallbackValue();
            m_needToUpdateShaderVariant = false;
            InvalidateSRG();
        }

        void SMAABasePass::InvalidateShaderVariant()
        {
            m_needToUpdateShaderVariant = true;
        }

        void SMAABasePass::InvalidateSRG()
        {
            m_needToUpdateSRG = true;
        }

        AZ::Vector4 SMAABasePass::CalculateRenderTargetMetrics(const RPI::PassAttachment* attachment)
        {
            AZ_Assert(attachment != nullptr, "Null PassAttachment pointer have been passed in SMAABasePass::CalculateRenderTargetMetrics().");

            const RPI::PassAttachmentBinding* sizeSource = attachment->m_sizeSource;
            AZ_Assert(sizeSource != nullptr, "Binding sizeSource of attachment is null.");
            AZ_Assert(sizeSource->m_attachment != nullptr, "Attachment of sizeSource is null.");
            AZ::RHI::Size size = sizeSource->m_attachment->m_descriptor.m_image.m_size;

            return AZ::Vector4(
                1.0f / static_cast<float>(size.m_width),
                1.0f / static_cast<float>(size.m_height),
                static_cast<float>(size.m_width),
                static_cast<float>(size.m_height));
        }

    }   // namespace Render
}   // namespace AZ