Asset.h

#pragma once
#include "Engine/Logging/TracyMacros.hpp"
#include <optional>
#include <cstdint>
#include <entt/entt.hpp>
#include <future>
#include <memory>
#include <string>
#include <utility>
#include <unordered_map>
#include <Engine/Ecs/RegistryManager.h>

namespace Asset
{
    enum CpuLoadingState : uint8_t {
        UNLOADED = 0,
        REQUESTED_LOAD = 1,
        LOADING = 2,
        LOADED = 3,
        FAILED = 4
    };
 
    using LoadState = CpuLoadingState;  
 
    template <typename Key, typename AssetClass>
    class AssetManager;
 
    template <typename Key, typename AssetClass>
    class Ref;

    template <typename Key, typename AssetClass>
    struct LoadResult
    {
        Ref<Key, AssetClass> ref;

        AssetClass* asset = nullptr;

        bool success = false;

        std::string error;

        uint64_t generation = 0;
    };

    template <typename Key, typename AssetClass>
    struct Record
    {
        using Result = LoadResult<Key, AssetClass>;  

        explicit Record(Key key) : key(std::move(key))
        {
            resetFuture();
        }

        void resetFuture()
        {
            readyPromise = std::make_shared<std::promise<Result>>();
            readyFuture = readyPromise->get_future().share();
            readyPromiseResolved = false;
        }

        Key key;

        std::unique_ptr<AssetClass> asset;

        LoadState state = UNLOADED;

        uint64_t generation = 0;

        std::string lastError;

        std::shared_ptr<std::promise<Result>> readyPromise;

        std::shared_future<Result> readyFuture;

        bool readyPromiseResolved = false;
    };

    template <typename Key, typename AssetClass>
    class Ref
    {
    public:
        using RecordType = Record<Key, AssetClass>;  
        using Result = LoadResult<Key, AssetClass>;  

        Ref() = default;

        [[nodiscard]] LoadState getState() const
        {
            const auto record = record_.lock();
            return record ? record->state : UNLOADED;
        }

        [[nodiscard]] bool isLoaded() const
        {
            return getState() == LOADED;
        }

        [[nodiscard]] AssetClass* get() const
        {
            const auto record = record_.lock();
            return record && record->state == LOADED ? record->asset.get() : nullptr;
        }

        [[nodiscard]] std::shared_ptr<RecordType> sharedRecord() const
        {
            return record_.lock();
        }

        [[nodiscard]] uint64_t getGeneration() const
        {
            const auto record = record_.lock();
            return record ? record->generation : 0;
        }

        [[nodiscard]] const std::string& getError() const
        {
            static const std::string empty;
            const auto record = record_.lock();
            return record ? record->lastError : empty;
        }

        [[nodiscard]] std::shared_future<Result> readyFuture() const
        {
            const auto record = record_.lock();
            return record ? record->readyFuture : expiredFuture();
        }

        [[nodiscard]] explicit operator bool() const
        {
            return !record_.expired();
        }
 
    private:
        friend class AssetManager<Key, AssetClass>;
 
        explicit Ref(std::shared_ptr<RecordType> record) : record_(std::move(record)) {}
 
        static std::shared_future<Result> expiredFuture()
        {
            static std::shared_future<Result> future = []
            {
                std::promise<Result> promise;
                Result result;
                result.error = "Asset record expired";
                promise.set_value(std::move(result));
                return promise.get_future().share();
            }();
            return future;
        }
 
        std::weak_ptr<RecordType> record_;
    };

    class AssetBase
    {
    public:
        explicit AssetBase(bool initializeDefaultEntity = false) :
            loadingState( initializeDefaultEntity ? LOADING : UNLOADED )
          , mainRegistry( Engine::Ecs::RegistryManager::get() )
          , data( initializeDefaultEntity ? mainRegistry.create() : entt::null )
        {
        };

        virtual ~AssetBase();

        [[nodiscard]] CpuLoadingState getLoadingState();

        void requestLoad();

        virtual void unload();

        entt::entity getEntity() const;
 
    protected:
        virtual void updateLoadingState()
        {
        }
        CpuLoadingState loadingState = UNLOADED;  
 
        entt::registry& mainRegistry;  
        entt::entity data;             
    };

    template <typename Key, typename AssetClass>
    class AssetManager {
    public:
        using RefType = Ref<Key, AssetClass>;          
        using RecordType = Record<Key, AssetClass>;    
        using ResultType = LoadResult<Key, AssetClass>;  

        AssetManager() = default;

        virtual ~AssetManager() = default;

        template <typename... Args>
        void declare( Key path, Args &&... args );

        void add(Key path, AssetClass* asset);

        RefType getOrCreateRef(const Key& path);

        [[nodiscard]] RefType getRef(const Key& path) const;

        void setAssetLoadingState(const Key& path, LoadState state);

        void resolve(const Key& path, std::unique_ptr<AssetClass> asset);

        void fail(const Key& path, std::string error);

        void unloadAssetPayload(const Key& path);

        bool exists(Key path);

        bool isDeclared(Key path);

        CpuLoadingState getAssetLoadingState(const Key& path);
        std::optional<AssetClass*> getAsset(const Key& path);

        virtual void clear();

        template<typename Func>
        void forEachAsset(Func&& func) const {
            for (const auto& [key, record] : records) {
                if (record && record->asset != nullptr) {
                    func(key, record->asset.get());
                }
            }
        }
 
    private:
        std::shared_ptr<RecordType> getOrCreateRecord(const Key& path);
        [[nodiscard]] std::shared_ptr<RecordType> getRecord(const Key& path) const;
        void completeRecord(const std::shared_ptr<RecordType>& record, bool success, std::string error = {});
        void syncRecordStateFromPayload(const std::shared_ptr<RecordType>& record) const;
 
        std::unordered_map<Key, std::shared_ptr<RecordType>> records;
    };
 
    template<typename Key, typename AssetClass>
    template<typename... Args>
    void AssetManager<Key, AssetClass>::declare(Key path, Args&&... args) {
        auto record = getOrCreateRecord(path);
        record->asset = std::make_unique<AssetClass>(std::forward<Args>(args)...);
        ++record->generation;
        record->resetFuture();
        syncRecordStateFromPayload(record);
        if (record->state == LOADED) {
            completeRecord(record, true);
        }
    }
 
    template<typename Key, typename AssetClass>
    void AssetManager<Key, AssetClass>::add(Key path, AssetClass *asset) {
        TRACY_ZONE_SCOPED_NAMED( "Added Asset to Asset Manager" );
        assert(asset != nullptr); // use declare for that
        auto record = getOrCreateRecord(path);
        record->asset.reset(asset);
        ++record->generation;
        record->resetFuture();
        syncRecordStateFromPayload(record);
        if (record->state == LOADED) {
            completeRecord(record, true);
        }
    }
 
    template<typename Key, typename AssetClass>
    typename AssetManager<Key, AssetClass>::RefType AssetManager<Key, AssetClass>::getOrCreateRef(const Key& path) {
        return RefType(getOrCreateRecord(path));
    }
 
    template<typename Key, typename AssetClass>
    typename AssetManager<Key, AssetClass>::RefType AssetManager<Key, AssetClass>::getRef(const Key& path) const {
        return RefType(getRecord(path));
    }
 
    template<typename Key, typename AssetClass>
    void AssetManager<Key, AssetClass>::setAssetLoadingState(const Key& path, LoadState state) {
        auto record = getOrCreateRecord(path);
 
        if ((state == REQUESTED_LOAD || state == LOADING) &&
            (record->state == UNLOADED || record->state == LOADED || record->state == FAILED))
        {
            ++record->generation;
            record->resetFuture();
            record->lastError.clear();
            if (state != LOADED) {
                record->asset.reset();
            }
        }
 
        record->state = state;
        if (state == LOADED) {
            completeRecord(record, true);
        }
    }
 
    template<typename Key, typename AssetClass>
    void AssetManager<Key, AssetClass>::resolve(const Key& path, std::unique_ptr<AssetClass> asset) {
        auto record = getOrCreateRecord(path);
        record->asset = std::move(asset);
        record->state = LOADED;
        record->lastError.clear();
        completeRecord(record, true);
    }
 
    template<typename Key, typename AssetClass>
    void AssetManager<Key, AssetClass>::fail(const Key& path, std::string error) {
        auto record = getOrCreateRecord(path);
        record->asset.reset();
        record->state = FAILED;
        record->lastError = std::move(error);
        completeRecord(record, false, record->lastError);
    }
 
    template<typename Key, typename AssetClass>
    void AssetManager<Key, AssetClass>::unloadAssetPayload(const Key& path) {
        auto record = getRecord(path);
        if (!record) return;
        record->asset.reset();
        record->state = UNLOADED;
        ++record->generation;
        record->resetFuture();
        record->lastError.clear();
    }
 
    template<typename Key, typename AssetClass>
    bool AssetManager<Key, AssetClass>::exists(Key path) {
        const auto record = getRecord(path);
        return record && record->asset != nullptr;
    }
 
    template <typename Key, typename AssetClass>
    bool AssetManager<Key, AssetClass>::isDeclared( Key path )
    {
        return records.find(path) != records.end();
    }
 
    template <typename Key, typename AssetClass>
    CpuLoadingState AssetManager<Key, AssetClass>::getAssetLoadingState(const Key& path) {
        const auto record = getRecord(path);
        if (!record) return UNLOADED;
        syncRecordStateFromPayload(record);
        return record->state;
    }
 
    template<typename Key, typename AssetClass>
    std::optional<AssetClass *> AssetManager<Key, AssetClass>::getAsset(const Key& path) {
        TRACY_ZONE_SCOPED_NAMED("Retrieve Asset");
        const auto record = getRecord(path);
        if (!record || record->asset == nullptr) return std::nullopt;
        syncRecordStateFromPayload(record);
        return record->asset.get();
    }
 
    template <typename Key, typename AssetClass>
    void AssetManager<Key, AssetClass>::clear()
    {
        TRACY_ZONE_SCOPED_NAMED( "Unloading data" );
        for (auto& [key, record] : records)
        {
            TRACY_ZONE_SCOPED_NAMED("Deleting loaded data");
            if (record) {
                record->asset.reset();
                record->state = UNLOADED;
                ++record->generation;
                record->resetFuture();
                record->lastError.clear();
            }
        }
        records.clear();
    }
 
    template<typename Key, typename AssetClass>
    std::shared_ptr<typename AssetManager<Key, AssetClass>::RecordType> AssetManager<Key, AssetClass>::getOrCreateRecord(const Key& path) {
        auto it = records.find(path);
        if (it != records.end()) {
            return it->second;
        }
 
        auto record = std::make_shared<RecordType>(path);
        records.emplace(path, record);
        return record;
    }
 
    template<typename Key, typename AssetClass>
    std::shared_ptr<typename AssetManager<Key, AssetClass>::RecordType> AssetManager<Key, AssetClass>::getRecord(const Key& path) const {
        const auto it = records.find(path);
        return it != records.end() ? it->second : nullptr;
    }
 
    template<typename Key, typename AssetClass>
    void AssetManager<Key, AssetClass>::completeRecord(const std::shared_ptr<RecordType>& record, bool success, std::string error) {
        if (!record || record->readyPromiseResolved || !record->readyPromise) return;
 
        ResultType result;
        result.ref = RefType(record);
        result.asset = success ? record->asset.get() : nullptr;
        result.success = success;
        result.error = std::move(error);
        result.generation = record->generation;
 
        record->readyPromise->set_value(std::move(result));
        record->readyPromiseResolved = true;
    }
 
    template<typename Key, typename AssetClass>
    void AssetManager<Key, AssetClass>::syncRecordStateFromPayload(const std::shared_ptr<RecordType>& record) const {
        if (!record || record->asset == nullptr) return;
        if (record->state == FAILED) return;
 
        const LoadState payloadState = record->asset->getLoadingState();
        if (payloadState != record->state) {
            record->state = payloadState;
        }
        if (record->state == LOADED) {
            const_cast<AssetManager*>(this)->completeRecord(record, true);
        }
    }
}  // namespace Asset