/*
* 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 <AzslShaderBuilderSystemComponent.h>

#include <AssetBuilderSDK/AssetBuilderSDK.h>
#include <AssetBuilderSDK/AssetBuilderBusses.h>

#include <Atom/RHI.Edit/ShaderPlatformInterface.h>
#include <Atom/RHI.Edit/Utils.h>

#include <Atom/RPI.Reflect/Asset/AssetHandler.h>
#include <Atom/RPI.Reflect/Shader/ShaderAsset.h>
#include <Atom/RPI.Edit/Shader/ShaderSourceData.h>
#include <Atom/RPI.Edit/Shader/ShaderVariantListSourceData.h>
#include <Atom/RPI.Reflect/Shader/ShaderResourceGroupAsset.h>

#include <AzToolsFramework/ToolsComponents/ToolsAssetCatalogBus.h>

#include <AzCore/Serialization/EditContext.h>
#include <AzCore/Serialization/SerializeContext.h>

#include <CommonFiles/Preprocessor.h>
#include <CommonFiles/GlobalBuildOptions.h>
#include <Editor/AtomShaderCapabilitiesConfigFile.h>

namespace AZ
{
    namespace ShaderBuilder
    {
        void AzslShaderBuilderSystemComponent::Reflect(ReflectContext* context)
        {
            if (SerializeContext* serialize = azrtti_cast<SerializeContext*>(context))
            {
                serialize->Class<AzslShaderBuilderSystemComponent, Component>()
                    ->Version(0)
                    ->Attribute(Edit::Attributes::SystemComponentTags, AZStd::vector<Crc32>({ AssetBuilderSDK::ComponentTags::AssetBuilder }))
                    ;
            }

            PreprocessorOptions::Reflect(context);
            RHI::ShaderCompilerProfiling::Reflect(context);
            AtomShaderConfig::CapabilitiesConfigFile::Reflect(context);
            GlobalBuildOptions::Reflect(context);
            RHI::ShaderCompilerArguments::Reflect(context);
        }

        void AzslShaderBuilderSystemComponent::GetProvidedServices(ComponentDescriptor::DependencyArrayType& provided)
        {
            provided.push_back(AZ_CRC("AzslShaderBuilderService", 0x09315a40));
        }

        void AzslShaderBuilderSystemComponent::GetIncompatibleServices(ComponentDescriptor::DependencyArrayType& incompatible)
        {
            incompatible.push_back(AZ_CRC("AzslShaderBuilderService", 0x09315a40));
        }

        void AzslShaderBuilderSystemComponent::GetRequiredServices(ComponentDescriptor::DependencyArrayType& required)
        {
            (void)required;
        }

        void AzslShaderBuilderSystemComponent::GetDependentServices(ComponentDescriptor::DependencyArrayType& dependent)
        {
            dependent.push_back(AZ_CRC("AssetCatalogService", 0xc68ffc57));
        }

        void AzslShaderBuilderSystemComponent::Init()
        {
        }

        void AzslShaderBuilderSystemComponent::Activate()
        {
            RHI::ShaderPlatformInterfaceRegisterBus::Handler::BusConnect();
            ShaderPlatformInterfaceRequestBus::Handler::BusConnect();

            // Register AZSL's compilation products Builder
            AssetBuilderSDK::AssetBuilderDesc azslBuilderDescriptor;
            azslBuilderDescriptor.m_name = "AZSL Builder";
            azslBuilderDescriptor.m_version = 7; // LKG Merge
            // register all extensions thay may carry azsl code. header. main shader. or SRG
            azslBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern(AZStd::string::format("*.%s", RPI::ShaderSourceData::Extension), AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
            azslBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern("*.azsl", AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
            azslBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern("*.azsli", AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
            azslBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern(AZStd::string::format("*.%s", SrgLayoutBuilder::MergedPartialSrgsExtension), AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
            azslBuilderDescriptor.m_busId = AzslBuilder::GetUUID();
            azslBuilderDescriptor.m_createJobFunction = AZStd::bind(&AzslBuilder::CreateJobs, &m_azslBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);
            azslBuilderDescriptor.m_processJobFunction = AZStd::bind(&AzslBuilder::ProcessJob, &m_azslBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);

            m_azslBuilder.BusConnect(azslBuilderDescriptor.m_busId);
            AssetBuilderSDK::AssetBuilderBus::Broadcast(&AssetBuilderSDK::AssetBuilderBus::Handler::RegisterBuilderInformation, azslBuilderDescriptor);

            // Register Shader Resource Group Layout Builder
            AssetBuilderSDK::AssetBuilderDesc srgLayoutBuilderDescriptor;
            srgLayoutBuilderDescriptor.m_name = "Shader Resource Group Layout Builder";
            srgLayoutBuilderDescriptor.m_version = 53; // ATOM-15196

            srgLayoutBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern("*.azsl", AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
            srgLayoutBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern("*.azsli", AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
            srgLayoutBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern(AZStd::string::format("*.%s", SrgLayoutBuilder::MergedPartialSrgsExtension), AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
            srgLayoutBuilderDescriptor.m_busId = SrgLayoutBuilder::GetUUID();
            srgLayoutBuilderDescriptor.m_createJobFunction = AZStd::bind(&SrgLayoutBuilder::CreateJobs, &m_srgLayoutBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);
            srgLayoutBuilderDescriptor.m_processJobFunction = AZStd::bind(&SrgLayoutBuilder::ProcessJob, &m_srgLayoutBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);

            m_srgLayoutBuilder.BusConnect(srgLayoutBuilderDescriptor.m_busId);
            AssetBuilderSDK::AssetBuilderBus::Broadcast(&AssetBuilderSDK::AssetBuilderBus::Handler::RegisterBuilderInformation, srgLayoutBuilderDescriptor);
            m_srgLayoutBuilder.Activate();

            // Register Shader Asset Builder
            AssetBuilderSDK::AssetBuilderDesc shaderAssetBuilderDescriptor;
            shaderAssetBuilderDescriptor.m_name = "Shader Asset Builder";
            shaderAssetBuilderDescriptor.m_version = 97; // ATOM-15196
            // .shader file changes trigger rebuilds
            shaderAssetBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern( AZStd::string::format("*.%s", RPI::ShaderSourceData::Extension), AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
            shaderAssetBuilderDescriptor.m_busId = azrtti_typeid<ShaderAssetBuilder>();
            shaderAssetBuilderDescriptor.m_createJobFunction = AZStd::bind(&ShaderAssetBuilder::CreateJobs, &m_shaderAssetBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);
            shaderAssetBuilderDescriptor.m_processJobFunction = AZStd::bind(&ShaderAssetBuilder::ProcessJob, &m_shaderAssetBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);

            m_shaderAssetBuilder.BusConnect(shaderAssetBuilderDescriptor.m_busId);
            AssetBuilderSDK::AssetBuilderBus::Broadcast(&AssetBuilderSDK::AssetBuilderBus::Handler::RegisterBuilderInformation, shaderAssetBuilderDescriptor);

            // Register Shader Variant Asset Builder
            AssetBuilderSDK::AssetBuilderDesc shaderVariantAssetBuilderDescriptor;
            shaderVariantAssetBuilderDescriptor.m_name = "Shader Variant Asset Builder";
            // Both "Shader Variant Asset Builder" and "Shader Asset Builder" produce ShaderVariantAsset products. If you update
            // ShaderVariantAsset you will need to update BOTH version numbers, not just "Shader Variant Asset Builder".
            shaderVariantAssetBuilderDescriptor.m_version = 18; // ATOM-15196
            shaderVariantAssetBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern(AZStd::string::format("*.%s", RPI::ShaderVariantListSourceData::Extension), AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
            shaderVariantAssetBuilderDescriptor.m_busId = azrtti_typeid<ShaderVariantAssetBuilder>();
            shaderVariantAssetBuilderDescriptor.m_createJobFunction = AZStd::bind(&ShaderVariantAssetBuilder::CreateJobs, &m_shaderVariantAssetBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);
            shaderVariantAssetBuilderDescriptor.m_processJobFunction = AZStd::bind(&ShaderVariantAssetBuilder::ProcessJob, &m_shaderVariantAssetBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);

            m_shaderVariantAssetBuilder.BusConnect(shaderVariantAssetBuilderDescriptor.m_busId);
            AssetBuilderSDK::AssetBuilderBus::Broadcast(&AssetBuilderSDK::AssetBuilderBus::Handler::RegisterBuilderInformation, shaderVariantAssetBuilderDescriptor);

            // Register Precompiled Shader Builder
            AssetBuilderSDK::AssetBuilderDesc precompiledShaderBuilderDescriptor;
            precompiledShaderBuilderDescriptor.m_name = "Precompiled Shader Builder";
            precompiledShaderBuilderDescriptor.m_version = 7; // ATOM-14780
            precompiledShaderBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern(AZStd::string::format("*.%s", AZ::PrecompiledShaderBuilder::Extension), AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
            precompiledShaderBuilderDescriptor.m_busId = azrtti_typeid<PrecompiledShaderBuilder>();
            precompiledShaderBuilderDescriptor.m_createJobFunction = AZStd::bind(&PrecompiledShaderBuilder::CreateJobs, &m_precompiledShaderBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);
            precompiledShaderBuilderDescriptor.m_processJobFunction = AZStd::bind(&PrecompiledShaderBuilder::ProcessJob, &m_precompiledShaderBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);

            m_precompiledShaderBuilder.BusConnect(precompiledShaderBuilderDescriptor.m_busId);
            AssetBuilderSDK::AssetBuilderBus::Broadcast(&AssetBuilderSDK::AssetBuilderBus::Handler::RegisterBuilderInformation, precompiledShaderBuilderDescriptor);
        }

        void AzslShaderBuilderSystemComponent::Deactivate()
        {
            m_shaderAssetBuilder.BusDisconnect();
            m_srgLayoutBuilder.BusDisconnect();
            m_shaderVariantAssetBuilder.BusDisconnect();
            m_precompiledShaderBuilder.BusDisconnect();

            RHI::ShaderPlatformInterfaceRegisterBus::Handler::BusDisconnect();
            ShaderPlatformInterfaceRequestBus::Handler::BusDisconnect();
        }

        void AzslShaderBuilderSystemComponent::RegisterShaderPlatformHandler(RHI::ShaderPlatformInterface* shaderPlatformInterface)
        {
            m_shaderPlatformInterfaces[shaderPlatformInterface->GetAPIType()] = shaderPlatformInterface;
        }

        void AzslShaderBuilderSystemComponent::UnregisterShaderPlatformHandler(RHI::ShaderPlatformInterface* shaderPlatformInterface)
        {
            m_shaderPlatformInterfaces.erase(shaderPlatformInterface->GetAPIType());
        }

        AZStd::vector<RHI::ShaderPlatformInterface*> AzslShaderBuilderSystemComponent::GetShaderPlatformInterface(const AssetBuilderSDK::PlatformInfo& platformInfo)
        {
            //Used to validate that each ShaderPlatformInterface returns a unique index.
            AZStd::unordered_map<uint32_t, Name> apiUniqueIndexToName;


            AZStd::vector<RHI::ShaderPlatformInterface*> results;
            results.reserve(platformInfo.m_tags.size());
            for (const AZStd::string& tag : platformInfo.m_tags)
            {
                // Use the platform tags to find the proper ShaderPlatformInterface
                auto findIt = m_shaderPlatformInterfaces.find(RHI::APIType(tag.c_str()));
                if (findIt != m_shaderPlatformInterfaces.end())
                {
                    RHI::ShaderPlatformInterface* rhiApi = findIt->second;
                    const uint32_t uniqueIndex = rhiApi->GetAPIUniqueIndex();
                    auto uniqueIt = apiUniqueIndexToName.find(uniqueIndex);
                    if (uniqueIt != apiUniqueIndexToName.end())
                    {
                        AZ_Assert(false, "The ShaderPlatformInterface with name [%s] is providing a unique api index [%u] which was already provided by the ShaderPlatformInterface [%s]",
                            rhiApi->GetAPIName().GetCStr(), uniqueIndex, uniqueIt->second.GetCStr());
                        continue;
                    }
                    AZ_Assert(uniqueIndex <= RHI::Limits::APIType::PerPlatformApiUniqueIndexMax, "The api index [%u] from ShaderPlatformInterface [%s] is invalid", uniqueIndex, rhiApi->GetAPIName().GetCStr());
                    apiUniqueIndexToName[uniqueIndex] = rhiApi->GetAPIName();

                    results.push_back(rhiApi);
                }
            }
            return results;
        }

    } // namespace ShaderBuilder
} // namespace AZ