AssetManager.h

#pragma once
#include "DescriptorIndexAllocator.h"
#include "Engine/Animation/AnimationAsset.h"
#include "Engine/Animation/AnimationAssetRef.h"
#include "Engine/Animation/SkinAsset.h"
#include "Engine/Animation/SkinAssetRef.h"
#include "Engine/Audio/AudioAsset.h"
#include "Engine/Audio/AudioAssetRef.h"
#include "Engine/Core/Settings.h"
#include "Engine/Ecs/FrameProcessing.h"
#include "Engine/Logging/TracyMacros.hpp"
#include "Engine/Material/MaterialAsset.h"
#include "Engine/Material/MaterialAssetRef.h"
#include "Engine/Texture/CubemapTexture.h"
#include "Engine/Texture/Texture.h"
#include "Engine/Texture/TextureAsset.h"
#include "Engine/Texture/TextureAssetRef.h"
#include "Engine/Texture/TextureHandleRegistry.h"
#include "Engine/Texture/RuntimeTextureDescriptorLease.h"
#include "MeshAsset.h"
#include "MeshAssetRef.h"
 
#include <Engine/Math/BidirectionalMap.h>
#include <Engine/Texture/TextureLoadData.h>
#include <filesystem>
#include <memory>
#include <optional>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <unordered_set>
 
namespace Engine::Entities
{
    class GltfSpawner;
}  // namespace Engine::Entities
 
namespace Engine::Rendering
{
    class Renderer;
}  // namespace Engine::Rendering
 
class NamedThreadPool;
 
namespace Engine::Assets
{
    class Material;
    class Mesh;
    class ModelAssetManager;
    class MeshAssetManager;
    class RuntimeTexture;
    class Texture;
}  // namespace Engine::Assets
 
namespace Engine::Core
{
    class ApplicationContext;
    class RenderingDataManager;
    class TextureManager;

    struct PathHasher
    {
        std::size_t operator()( const std::filesystem::path & p ) const noexcept
        {
            return std::hash<std::string>{}( p.string() );  // Hash the string representation
        }
    };

    struct PackedVertex
    {
        glm::vec3 position;        
        uint32_t  packedNormal;    
        uint32_t  packedColor;     
        uint32_t  packedTexCoord;  
    };

    inline uint32_t packNormalA2B10G10R10_SNORM( const glm::vec3 & n )
    {
        // Clamp input just in case
        float x = std::max( -1.0f, std::min( 1.0f, n.x ) );
        float y = std::max( -1.0f, std::min( 1.0f, n.y ) );
        float z = std::max( -1.0f, std::min( 1.0f, n.z ) );
 
        // Scale to signed 10-bit range [-511, 511]
        int32_t ix = static_cast<int32_t>( std::round( x * 511.0f ) );
        int32_t iy = static_cast<int32_t>( std::round( y * 511.0f ) );
        int32_t iz = static_cast<int32_t>( std::round( z * 511.0f ) );
 
        // Pack into uint32_t (masking handles two's complement correctly)
        // Layout: | 2 unused | 10 bits B | 10 bits G | 10 bits R |
        return ( static_cast<uint32_t>( iz & 0x3FF ) << 20 ) | ( static_cast<uint32_t>( iy & 0x3FF ) << 10 ) |
               ( static_cast<uint32_t>( ix & 0x3FF ) );
        // Alpha bits (highest 2) remain 0
    }

    inline uint32_t packColorR8G8B8A8_UNORM( const glm::vec3 & c )
    {
        // Clamp input
        float r = std::max( 0.0f, std::min( 1.0f, c.x ) );
        float g = std::max( 0.0f, std::min( 1.0f, c.y ) );
        float b = std::max( 0.0f, std::min( 1.0f, c.z ) );
 
        // Scale to [0, 255]
        uint32_t ur = static_cast<uint32_t>( std::round( r * 255.0f ) );
        uint32_t ug = static_cast<uint32_t>( std::round( g * 255.0f ) );
        uint32_t ub = static_cast<uint32_t>( std::round( b * 255.0f ) );
        uint32_t ua = 255;  // Fully opaque alpha
 
        // Pack (assuming typical RGBA little-endian layout)
        return ( ua << 24 ) | ( ub << 16 ) | ( ug << 8 ) | ur;
    }

    inline uint32_t packTexCoordR16G16_SFLOAT( const glm::vec2 & tc )
    {
        // Use glm's packing function (requires <glm/gtc/packing.hpp>)
        // It converts two floats to two 16-bit half-floats and packs them.
        return glm::packHalf2x16( tc );
    }

    struct TextureStorage
    {
        uint32_t index = 0u;  

        std::shared_ptr<Assets::Texture> texture = nullptr;

        template <typename T>
        bool operator>( const T & other )
        {
            return index > other.index;
        }

        template <typename T>
        bool operator<( const T & other )
        {
            return index < other.index;
        }
    };

    class AssetManager
    {
        friend ApplicationContext;
        friend RenderingDataManager;
        friend Core::TextureManager;
        friend Entities::GltfSpawner;
 
    public:
        using AudioAssetRef     = Assets::AudioAssetRef;
        using AnimationAssetRef = Assets::AnimationAssetRef;
        using MaterialAssetRef  = Assets::MaterialAssetRef;
        using MeshAssetRef      = Assets::MeshAssetRef;
        using SkinAssetRef      = Assets::SkinAssetRef;
        using TextureAssetRef   = Assets::TextureAssetRef;

        template <typename AssetClass, typename Key>
        [[nodiscard]] auto getAsset( const Key & asset )
        {
            if constexpr ( std::is_same_v<AssetClass, Assets::Mesh> )
            {
                static_assert(
                    std::is_same_v<std::decay_t<Key>, Asset::MeshPath>,
                    "Mesh assets are keyed by Asset::MeshPath"
                );
 
                MeshAssetRef ref = meshAssetManager->getOrCreateRef( asset );
                (void)getMeshAsset( asset );
                return ref;
            }
            else if constexpr ( std::is_same_v<AssetClass, Assets::Audio> )
            {
                static_assert(
                    std::is_same_v<std::decay_t<Key>, Asset::Path> ||
                        std::is_same_v<std::decay_t<Key>, std::filesystem::path>,
                    "Audio assets are keyed by Asset::Path"
                );
 
                return loadAudioAsset( Asset::Path( asset ) );
            }
            else if constexpr ( std::is_same_v<AssetClass, Assets::Animation> )
            {
                static_assert(
                    std::is_same_v<std::decay_t<Key>, Asset::MeshPath>,
                    "Animation assets are keyed by Asset::MeshPath"
                );
 
                AnimationAssetRef ref = animationAssetManager_->getOrCreateRef( asset );
                if ( animationAssetManager_->getAssetLoadingState( asset ) == Asset::UNLOADED )
                {
                    loadEcsModel( asset.getFilePath() );
                }
                return ref;
            }
            else if constexpr ( std::is_same_v<AssetClass, Assets::Material> )
            {
                static_assert(
                    std::is_same_v<std::decay_t<Key>, Asset::MeshPath>,
                    "Material assets are keyed by Asset::MeshPath"
                );
 
                return materialAssetManager->getOrCreateRef( asset );
            }
            else if constexpr ( std::is_same_v<AssetClass, Assets::Skin> )
            {
                static_assert(
                    std::is_same_v<std::decay_t<Key>, Asset::MeshPath>,
                    "Skin assets are keyed by Asset::MeshPath"
                );
 
                SkinAssetRef ref = skinAssetManager_->getOrCreateRef( asset );
                if ( skinAssetManager_->getAssetLoadingState( asset ) == Asset::UNLOADED )
                {
                    loadEcsModel( asset.getFilePath() );
                }
                return ref;
            }
            else if constexpr ( std::is_same_v<AssetClass, Assets::Texture> )
            {
                static_assert(
                    std::is_same_v<std::decay_t<Key>, Asset::Path> ||
                        std::is_same_v<std::decay_t<Key>, std::filesystem::path>,
                    "Texture assets are keyed by Asset::Path"
                );
 
                const Asset::Path      pathKey( asset );
                TextureAssetRef        ref   = textureAssetManager->getOrCreateRef( pathKey );
                const Asset::LoadState state = textureAssetManager->getAssetLoadingState( pathKey );
                if ( state != Asset::LOADED && state != Asset::REQUESTED_LOAD && state != Asset::LOADING )
                {
                    loadEcsTexture( pathKey.getFilePath() );
                }
                return ref;
            }
            else
            {
                static_assert(
                    std::is_same_v<AssetClass, void>, "Unsupported asset type for AssetManager::getAsset<T>()"
                );
            }
        }

        AssetManager();

        AssetManager( ApplicationContext * context );

        ~AssetManager();

        void attachToRenderer( Rendering::Renderer * renderer );

        void setRenderingDataManager( RenderingDataManager * renderingDataManager );
 
    private:
        void loadEcsModel( const std::filesystem::path & path );

        void loadEcsTexture( const std::filesystem::path & path );
 
    public:
        Assets::CubemapTexture * loadEcsCubemap( const std::filesystem::path & path );

        [[nodiscard]] uint32_t getTextureDescriptorIndex( const std::filesystem::path & path );

        [[nodiscard]] uint32_t getTextureDescriptorIndex(
            const std::filesystem::path & path,
            Assets::TextureType          type
        );

        [[nodiscard]] Assets::TextureHandle getOrCreateTextureHandleWithSlot(
            const std::filesystem::path & path,
            Assets::TextureType          type
        );

        template <Assets::TextureType Type>
        [[nodiscard]] Assets::TypedTextureHandle<Type> getOrCreateTextureHandleWithSlot(
            const std::filesystem::path & path
        )
        {
            return Assets::TypedTextureHandle<Type>( getOrCreateTextureHandleWithSlot( path, Type ) );
        }

        uint32_t getImageCount() const;

        void cleanup();

        std::vector<VkDescriptorImageInfo> getTextureDescriptorInfos() const;

        Ecs::AssetImportPipeline & getAssetImportPipeline();

        Ecs::TextureAssetPipeline & getTextureAssetPipeline();

        Ecs::ModelAssetPipeline &   getMeshAssetPipeline();

        Ecs::AudioAssetPipeline &   getAudioAssetPipeline();

        void                        unloadAllData();

        void discardSceneLoadingData();

        AudioAssetRef loadAudioAsset( const Asset::Path & path );

        [[nodiscard]] Assets::Audio * getAudioAsset( const Asset::Path & path );

        [[nodiscard]] Assets::AudioAssetManager * getAudioAssetManager() const
        {
            return audioAssetManager_;
        }
 
    private:
        [[nodiscard]] Assets::Mesh * getMeshAsset( const Asset::MeshPath & asset );
 
    public:
        [[nodiscard]] Assets::MaterialAssetManager * getMaterialAssetManager() const
        {
            return materialAssetManager;
        }

        [[nodiscard]] Assets::ModelAssetManager * getModelAssetManager() const
        {
            return modelAssetManager;
        }

        [[nodiscard]] Assets::MeshAssetManager * getMeshAssetManager() const
        {
            return meshAssetManager;
        }

        [[nodiscard]] Assets::SkinAssetManager * getSkinAssetManager() const
        {
            return skinAssetManager_;
        }

        [[nodiscard]] Assets::AnimationAssetManager * getAnimationAssetManager() const
        {
            return animationAssetManager_;
        }

        [[nodiscard]] Assets::Textures::TextureHandleRegistry * getTextureHandleRegistry()
        {
            return &textureHandleRegistry_;
        }

        [[nodiscard]] Assets::RuntimeTextureDescriptorLease reserveRuntimeTextureDescriptor();
 
    private:
        friend class Assets::RuntimeTextureDescriptorLease;
 
        void releaseRuntimeTextureDescriptorIndex( uint32_t descriptorIndex );
 
        NamedThreadPool * threadPool_          = nullptr;
        NamedThreadPool * threadedCalculation_ = nullptr;
 
        Assets::TextureAssetManager *   textureAssetManager;
        Assets::MeshAssetManager *      meshAssetManager;
        Assets::MaterialAssetManager *  materialAssetManager;
        Assets::ModelAssetManager *     modelAssetManager;
        Assets::SkinAssetManager *      skinAssetManager_;
        Assets::AnimationAssetManager * animationAssetManager_;
 
        // IMPORTANT: textureHandleRegistry_ must be declared BEFORE assetImportPipeline
        // because the model pipeline's constructor receives 'this' and may access the registry.
        // C++ initializes members in declaration order, not initializer list order.
        Assets::Textures::TextureHandleRegistry
            textureHandleRegistry_;  
 
        Assets::AudioAssetManager * audioAssetManager_;
        Ecs::AssetImportPipeline   assetImportPipeline;
 
    public:
        NamedThreadPool * getCalculationPool() const
        {
            return threadedCalculation_;
        }

        NamedThreadPool * getAssetLoaderPool() const
        {
            return threadPool_;
        }
 
    private:
        Assets::Loaders::GltfLoader gltfLoader = Assets::Loaders::GltfLoader( threadPool_ );
 
        ApplicationContext *  context  = nullptr;
        Rendering::Renderer * renderer = nullptr;
 
        std::unordered_map<
            Assets::Textures::TextureIdentityKey,
            Assets::RuntimeTexture *,
            Assets::Textures::TextureIdentityKeyHash>
            textures = {};
        std::unordered_map<std::string, std::vector<Assets::Textures::TextureIdentityKey>> textureIdentitiesByPath_;
        std::unordered_map<std::filesystem::path, std::unique_ptr<Assets::CubemapTexture>, PathHasher>
                                                              cubemapTextures = {};
        BidirectionalMap<uint32_t, Assets::RuntimeTexture *> textureIndexMap;
 
        std::array<std::optional<Assets::RuntimeTexture>, MAX_TEXTURE_COUNT> textureData;
        DescriptorIndexAllocator                                             textureAllocator;
        std::unordered_set<uint32_t>                                         runtimeTextureDescriptorIndices_;
        std::unordered_set<uint32_t>                                         fileTextureDescriptorIndices_;
        uint32_t fallbackTextureDescriptorIndex_ = Assets::TextureHandle::INVALID_ID;

        bool doesTextureAlreadyExist( const std::filesystem::path & path );

        Assets::RuntimeTexture *
            registerTexture( const std::filesystem::path & path, Assets::RuntimeTexture texture );

        Assets::RuntimeTexture * registerTexture(
            const std::filesystem::path & path,
            Assets::TextureType          type,
            Assets::RuntimeTexture       texture
        );

        Assets::RuntimeTexture *
            registerTextureFromPrepared( const std::filesystem::path & path, Assets::RuntimeTexture texture );

        Assets::RuntimeTexture * registerTextureFromPrepared(
            const std::filesystem::path & path,
            Assets::TextureType          type,
            Assets::RuntimeTexture       texture
        );

        Assets::RuntimeTexture *
            replaceTextureFromPrepared( const std::filesystem::path & path, Assets::RuntimeTexture texture );

        Assets::RuntimeTexture * replaceTextureFromPrepared(
            const std::filesystem::path & path,
            Assets::TextureType          type,
            Assets::RuntimeTexture       texture
        );

        void unregisterTexture( Assets::RuntimeTexture * texture );

        void unloadRuntimeTextures();

        void ensureFallbackTextureLocked();
 
        TRACY_LOCKABLE( std::mutex, textureMutex, "Texture mutex" );
 
        std::vector<TextureStorage>           texturesToUpload{};
        std::vector<Assets::RuntimeTexture *> texturesToCopyImageData;
    };
}  // namespace Engine::Core