text: /*
   * 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/RPI.Public/Image/AttachmentImage.h>
  #include <Atom/RPI.Public/Image/AttachmentImagePool.h>
  #include <Atom/RPI.Public/Image/ImageSystem.h>
  #include <Atom/RPI.Public/Image/StreamingImage.h>
  #include <Atom/RPI.Public/Image/StreamingImagePool.h>
  #include <Atom/RPI.Public/Image/DefaultStreamingImageController.h>
  
  #include <Atom/RPI.Reflect/Asset/AssetHandler.h>
  #include <Atom/RPI.Reflect/Image/AttachmentImageAssetCreator.h>
  #include <Atom/RPI.Reflect/Image/ImageMipChainAsset.h>
  #include <Atom/RPI.Reflect/Image/StreamingImageAssetHandler.h>
  #include <Atom/RPI.Reflect/Image/StreamingImagePoolAssetCreator.h>
  #include <Atom/RPI.Reflect/ResourcePoolAssetCreator.h>
  
  #include <Atom/RHI.Reflect/ImagePoolDescriptor.h>
  #include <Atom/RHI/ImagePool.h>
  
  #include <Atom/RHI/CpuProfiler.h>
  #include <Atom/RHI/RHISystemInterface.h>
  
  #include <AtomCore/Instance/InstanceDatabase.h>
  
  #include <AzCore/Debug/EventTrace.h>
  #include <AzCore/Interface/Interface.h>
  #include <AzCore/Interface/Interface.h>
  #include <AzCore/Asset/AssetManager.h>
  #include <AzCore/Math/Color.h>
  
  namespace AZ
  {
      namespace RPI
      {
          ImageSystemInterface* ImageSystemInterface::Get()
          {
              return Interface<ImageSystemInterface>::Get();
          }
  
          void ImageSystem::Reflect(AZ::ReflectContext* context)
          {
              ImageAsset::Reflect(context);
              ImageMipChainAsset::Reflect(context);
              ImageSystemDescriptor::Reflect(context);
              StreamingImageAsset::Reflect(context);
              StreamingImagePoolAsset::Reflect(context);
              StreamingImageControllerAsset::Reflect(context);
              DefaultStreamingImageControllerAsset::Reflect(context);
              AttachmentImageAsset::Reflect(context);
          }
  
          void ImageSystem::GetAssetHandlers(AssetHandlerPtrList& assetHandlers)
          {
              assetHandlers.emplace_back(MakeAssetHandler<AttachmentImageAssetHandler>());
              assetHandlers.emplace_back(MakeAssetHandler<ImageMipChainAssetHandler>());
              assetHandlers.emplace_back(MakeAssetHandler<StreamingImageAssetHandler>());
              assetHandlers.emplace_back(MakeAssetHandler<StreamingImagePoolAssetHandler>());
              assetHandlers.emplace_back(MakeAssetHandler<BuiltInAssetHandler>(
                  azrtti_typeid<DefaultStreamingImageControllerAsset>(),
                  []() { return aznew DefaultStreamingImageControllerAsset(); }));
          }
  
          void ImageSystem::Init(const ImageSystemDescriptor& desc)
          {
              RHI::Ptr<RHI::Device> device = RHI::RHISystemInterface::Get()->GetDevice();
  
              // Register attachment image instance database.
  
              {
                  Data::InstanceHandler<AttachmentImage> handler;
                  handler.m_createFunction = [](Data::AssetData* imageAsset)
                  {
                      return AttachmentImage::CreateInternal(*(azrtti_cast<AttachmentImageAsset*>(imageAsset)));
                  };
                  Data::InstanceDatabase<AttachmentImage>::Create(azrtti_typeid<AttachmentImageAsset>(), handler);
              }
  
              {
                  Data::InstanceHandler<AttachmentImagePool> handler;
                  handler.m_createFunction = [device](Data::AssetData* poolAsset)
                  {
                      return AttachmentImagePool::CreateInternal(*device, *(azrtti_cast<ResourcePoolAsset*>(poolAsset)));
                  };
                  Data::InstanceDatabase<AttachmentImagePool>::Create(azrtti_typeid<ResourcePoolAsset>(), handler);
              }
  
              // Register streaming image instance database.
  
              {
                  Data::InstanceHandler<StreamingImage> handler;
                  handler.m_createFunction = [](Data::AssetData* imageAsset)
                  {
                      return StreamingImage::CreateInternal(*(azrtti_cast<StreamingImageAsset*>(imageAsset)));
                  };
                  Data::InstanceDatabase<StreamingImage>::Create(azrtti_typeid<StreamingImageAsset>(), handler);
              }
  
              {
                  Data::InstanceHandler<StreamingImagePool> handler;
                  handler.m_createFunction = [this, device](Data::AssetData* poolAsset)
                  {
                      Data::Instance<StreamingImagePool> instance = StreamingImagePool::CreateInternal(*device, *(azrtti_cast<StreamingImagePoolAsset*>(poolAsset)));
                      if (instance)
                      {
                          m_activeStreamingPoolMutex.lock();
                          m_activeStreamingPools.emplace_back(instance.get());
                          m_activeStreamingPoolMutex.unlock();
                      }
                      return instance;
                  };
                  handler.m_deleteFunction = [this](StreamingImagePool* pool)
                  {
                      m_activeStreamingPoolMutex.lock();
                      auto findIt = AZStd::find(m_activeStreamingPools.begin(), m_activeStreamingPools.end(), pool);
                      AZ_Assert(findIt != m_activeStreamingPools.end(), "Pool must exist in the container.");
                      m_activeStreamingPools.erase(findIt);
                      m_activeStreamingPoolMutex.unlock();
  
                      delete pool;
                  };
                  Data::InstanceDatabase<StreamingImagePool>::Create(azrtti_typeid<StreamingImagePoolAsset>(), handler);
              }
  
              // Register streaming image controller instance database.
              {
                  Data::InstanceHandler<StreamingImageController> handler;
                  handler.m_createFunction = DefaultStreamingImageController::CreateInternal;
                  Data::InstanceDatabase<StreamingImageController>::Create(azrtti_typeid<StreamingImageControllerAsset>(), handler);
              }
  
              m_defaultStreamingImageControllerAsset = 
                  Data::AssetManager::Instance().CreateAsset<DefaultStreamingImageControllerAsset>(
                      DefaultStreamingImageControllerAsset::BuiltInAssetId, AZ::Data::AssetLoadBehavior::PreLoad);
  
              CreateDefaultResources(desc);
  
              Interface<ImageSystemInterface>::Register(this);
  
              m_initialized = true;
          }
  
          void ImageSystem::Shutdown()
          {
              if (!m_initialized)
              {
                  return;
              }
              Interface<ImageSystemInterface>::Unregister(this);
  
              m_defaultStreamingImageControllerAsset.Release();
              m_systemImages.clear();
              m_systemStreamingPool = nullptr;
              m_systemAttachmentPool = nullptr;
              m_assetStreamingPool = nullptr;
  
              Data::InstanceDatabase<AttachmentImage>::Destroy();
              Data::InstanceDatabase<AttachmentImagePool>::Destroy();
              Data::InstanceDatabase<StreamingImage>::Destroy();
              Data::InstanceDatabase<StreamingImagePool>::Destroy();
              Data::InstanceDatabase<StreamingImageController>::Destroy();
  
              m_activeStreamingPools.clear();
              m_initialized = false;
          }
  
          void ImageSystem::Update()
          {
              AZ_ATOM_PROFILE_FUNCTION("RPI", "ImageSystem: Update");
  
              AZStd::lock_guard<AZStd::mutex> lock(m_activeStreamingPoolMutex);
              for (StreamingImagePool* imagePool : m_activeStreamingPools)
              {
                  imagePool->Update();
              }
          }
  
          const Data::Instance<StreamingImagePool>& ImageSystem::GetSystemStreamingPool() const
          {
              return m_systemStreamingPool;
          }
          
          const Data::Instance<StreamingImagePool>& ImageSystem::GetStreamingPool() const
          {
              return m_assetStreamingPool;
          }
  
          const Data::Instance<AttachmentImagePool>& ImageSystem::GetSystemAttachmentPool() const
          {
              return m_systemAttachmentPool;
          }
  
          const Data::Instance<Image>& ImageSystem::GetSystemImage(SystemImage simpleImage) const
          {
              return m_systemImages[static_cast<size_t>(simpleImage)];
          }
  
          void ImageSystem::CreateDefaultResources(const ImageSystemDescriptor& desc)
          {
              struct SystemImageDescriptor
              {
                  SystemImageDescriptor(const Color color, const char* name)
                      : m_color{color}
                      , m_name{name}
                      , m_assetId{Uuid::CreateName(name)}
                  {}
  
                  Color m_color;
                  const char* m_name;
                  Data::AssetId m_assetId;
              };
  
              const SystemImageDescriptor systemImageDescriptors[static_cast<uint32_t>(SystemImage::Count)] =
              {
                  { Color(1.0f, 1.0f, 1.0f, 1.0f),  "Image_White" },
                  { Color(0.0f, 0.0f, 0.0f, 1.0f),  "Image_Black" },
                  { Color(0.5f, 0.5f, 0.5f, 1.0f),  "Image_Grey" },
                  { Color(1.0f, 0.0f, 1.0f, 1.0f),  "Image_Magenta" }
              };
  
              static_assert(AZ_ARRAY_SIZE(systemImageDescriptors) == static_cast<size_t>(SystemImage::Count), "System image arrays do not match.");
  
              struct SystemImagePoolDescriptor
              {
                  SystemImagePoolDescriptor(size_t budgetInBytes, const char* name)
                      : m_budgetInBytes{budgetInBytes}
                      , m_name{name}
                      , m_assetId{Uuid::CreateName(name)}
                  {}
  
                  size_t m_budgetInBytes;
                  const char* m_name;
                  Data::AssetId m_assetId;
              };
  
              const SystemImagePoolDescriptor systemStreamingPoolDescriptor{ desc.m_systemStreamingImagePoolSize, "ImageSystem::SystemStreamingImagePool" };
              const SystemImagePoolDescriptor systemAttachmentPoolDescriptor{desc.m_systemAttachmentImagePoolSize, "ImageSystem::AttachmentImagePool" };
              const SystemImagePoolDescriptor assetStreamingPoolDescriptor{ desc.m_assetStreamingImagePoolSize, "ImageSystem::AssetStreamingImagePool" };
  
              // Create the system streaming pool
              {
                  AZStd::unique_ptr<RHI::StreamingImagePoolDescriptor> imagePoolDescriptor = AZStd::make_unique<RHI::StreamingImagePoolDescriptor>();
                  imagePoolDescriptor->m_budgetInBytes = systemStreamingPoolDescriptor.m_budgetInBytes;
  
                  Data::Asset<StreamingImagePoolAsset> poolAsset;
  
                  StreamingImagePoolAssetCreator poolAssetCreator;
                  poolAssetCreator.Begin(systemStreamingPoolDescriptor.m_assetId);
                  poolAssetCreator.SetPoolDescriptor(AZStd::move(imagePoolDescriptor));
                  poolAssetCreator.SetControllerAsset(m_defaultStreamingImageControllerAsset);
                  poolAssetCreator.SetPoolName(systemStreamingPoolDescriptor.m_name);
                  [[maybe_unused]] const bool created = poolAssetCreator.End(poolAsset);
                  AZ_Assert(created, "Failed to build streaming image pool");
  
                  m_systemStreamingPool = StreamingImagePool::FindOrCreate(poolAsset);
              }
  
              // Create the streaming pool for streaming image from asset
              {
                  AZStd::unique_ptr<RHI::StreamingImagePoolDescriptor> imagePoolDescriptor = AZStd::make_unique<RHI::StreamingImagePoolDescriptor>();
                  imagePoolDescriptor->m_budgetInBytes = assetStreamingPoolDescriptor.m_budgetInBytes;
  
                  Data::Asset<StreamingImagePoolAsset> poolAsset;
  
                  StreamingImagePoolAssetCreator poolAssetCreator;
                  poolAssetCreator.Begin(assetStreamingPoolDescriptor.m_assetId);
                  poolAssetCreator.SetPoolDescriptor(AZStd::move(imagePoolDescriptor));
                  poolAssetCreator.SetControllerAsset(m_defaultStreamingImageControllerAsset);
                  poolAssetCreator.SetPoolName(assetStreamingPoolDescriptor.m_name);
                  [[maybe_unused]] const bool created = poolAssetCreator.End(poolAsset);
                  AZ_Assert(created, "Failed to build streaming image pool for assets");
  
                  m_assetStreamingPool = StreamingImagePool::FindOrCreate(poolAsset);
              }
              // Create the system attachment pool.
              {
                  AZStd::unique_ptr<RHI::ImagePoolDescriptor> imagePoolDescriptor = AZStd::make_unique<RHI::ImagePoolDescriptor>();
                  imagePoolDescriptor->m_budgetInBytes = systemAttachmentPoolDescriptor.m_budgetInBytes;
                  imagePoolDescriptor->m_bindFlags =
                      RHI::ImageBindFlags::ShaderRead | RHI::ImageBindFlags::ShaderWrite |
                      RHI::ImageBindFlags::Color | RHI::ImageBindFlags::DepthStencil |
                      RHI::ImageBindFlags::CopyRead | RHI::ImageBindFlags::CopyWrite;
  
                  Data::Asset<ResourcePoolAsset> poolAsset;
  
                  ResourcePoolAssetCreator poolAssetCreator;
                  poolAssetCreator.Begin(systemAttachmentPoolDescriptor.m_assetId);
                  poolAssetCreator.SetPoolDescriptor(AZStd::move(imagePoolDescriptor));
                  poolAssetCreator.SetPoolName(systemAttachmentPoolDescriptor.m_name);
                  [[maybe_unused]] const bool created = poolAssetCreator.End(poolAsset);
                  AZ_Assert(created, "Failed to build attachment image pool");
  
                  m_systemAttachmentPool = AttachmentImagePool::FindOrCreate(poolAsset);
              }
  
              // Create the set of system images.
              {
                  const size_t systemImageCount = static_cast<size_t>(SystemImage::Count);
                  m_systemImages.resize(systemImageCount);
  
                  for (size_t imageIndex = 0; imageIndex < systemImageCount; imageIndex++)
                  {
                      const uint32_t colorU32 = systemImageDescriptors[imageIndex].m_color.ToU32();
  
                      m_systemImages[imageIndex] = StreamingImage::CreateFromCpuData(
                          *m_systemStreamingPool,
                          RHI::ImageDimension::Image2D,
                          RHI::Size{ 1, 1, 1 },
                          RHI::Format::R8G8B8A8_UNORM_SRGB,
                          &colorU32,
                          sizeof(uint32_t));
                  }
              }
          }
      } // namespace RPI
  }// namespace AZ
  
  