XCEngine/engine/src/Core/Asset/ResourceManager.cpp

#include <XCEngine/Core/Asset/ResourceManager.h>
#include <XCEngine/Core/Asset/ResourceHandle.h>
#include <XCEngine/Core/Asset/ResourceTypes.h>
#include <XCEngine/Core/IO/ResourceFileSystem.h>
#include <XCEngine/Resources/Material/MaterialLoader.h>
#include <XCEngine/Resources/Mesh/MeshLoader.h>
#include <XCEngine/Resources/Shader/ShaderLoader.h>
#include <XCEngine/Resources/Texture/TextureLoader.h>
#include <XCEngine/Resources/UI/UIDocumentLoaders.h>
#include <XCEngine/Resources/Volume/VolumeFieldLoader.h>
#include <exception>

namespace XCEngine {
namespace Resources {

namespace {

std::string ToStdString(const Containers::String& value) {
    return std::string(value.CStr());
}

bool ShouldTraceResourcePath(const Containers::String& path) {
    const std::string text = ToStdString(path);
    return text.rfind("builtin://", 0) == 0 ||
           text.find("backpack") != std::string::npos ||
           text.find("New Material.mat") != std::string::npos;
}

Containers::String EncodeAssetRef(const AssetRef& assetRef) {
    if (!assetRef.IsValid()) {
        return Containers::String("<invalid>");
    }

    return assetRef.assetGuid.ToString() + "," +
           Containers::String(std::to_string(assetRef.localID).c_str()) + "," +
           Containers::String(std::to_string(static_cast<int>(assetRef.resourceType)).c_str());
}

template<typename TLoader>
void RegisterBuiltinLoader(ResourceManager& manager, TLoader& loader) {
    if (manager.GetLoader(loader.GetResourceType()) == nullptr) {
        manager.RegisterLoader(&loader);
    }
}

MaterialLoader g_materialLoader;
MeshLoader g_meshLoader;
ShaderLoader g_shaderLoader;
TextureLoader g_textureLoader;
UIViewLoader g_uiViewLoader;
UIThemeLoader g_uiThemeLoader;
UISchemaLoader g_uiSchemaLoader;
VolumeFieldLoader g_volumeFieldLoader;

} // namespace

ResourceManager& ResourceManager::Get() {
    static ResourceManager instance;
    return instance;
}

ResourceManager::ScopedDeferredSceneLoad::ScopedDeferredSceneLoad(ResourceManager& manager)
    : m_manager(&manager) {
    m_manager->BeginDeferredSceneLoad();
}

ResourceManager::ScopedDeferredSceneLoad::~ScopedDeferredSceneLoad() {
    if (m_manager != nullptr) {
        m_manager->EndDeferredSceneLoad();
    }
}

void ResourceManager::Initialize() {
    EnsureInitialized();
}

void ResourceManager::EnsureInitialized() {
    if (m_asyncLoader) {
        return;
    }

    std::lock_guard<std::mutex> initLock(m_initializeMutex);
    if (m_asyncLoader) {
        return;
    }

    Core::UniqueRef<AsyncLoader> asyncLoader = Core::MakeUnique<AsyncLoader>();
    asyncLoader->Initialize(2);

    RegisterBuiltinLoader(*this, g_materialLoader);
    RegisterBuiltinLoader(*this, g_meshLoader);
    RegisterBuiltinLoader(*this, g_shaderLoader);
    RegisterBuiltinLoader(*this, g_textureLoader);
    RegisterBuiltinLoader(*this, g_uiViewLoader);
    RegisterBuiltinLoader(*this, g_uiThemeLoader);
    RegisterBuiltinLoader(*this, g_uiSchemaLoader);
    RegisterBuiltinLoader(*this, g_volumeFieldLoader);
    m_assetImportService.Initialize();

    m_asyncLoader = std::move(asyncLoader);
}

void ResourceManager::Shutdown() {
    UnloadAll();
    if (m_asyncLoader) {
        m_asyncLoader->Shutdown();
        m_asyncLoader.reset();
    }

    m_assetImportService.Shutdown();
    ResourceFileSystem::Get().Shutdown();
    m_projectAssetIndex.ResetProjectRoot();

    std::lock_guard<std::mutex> inFlightLock(m_inFlightLoadsMutex);
    m_inFlightLoads.clear();
}

void ResourceManager::SetResourceRoot(const Containers::String& rootPath) {
    EnsureInitialized();
    m_resourceRoot = rootPath;
    if (!m_resourceRoot.Empty()) {
        ResourceFileSystem::Get().Initialize(rootPath);
        m_assetImportService.SetProjectRoot(rootPath);
        BootstrapProjectAssets();
    } else {
        m_assetImportService.SetProjectRoot(Containers::String());
        ResourceFileSystem::Get().Shutdown();
        m_projectAssetIndex.ResetProjectRoot();
    }
}

const Containers::String& ResourceManager::GetResourceRoot() const {
    return m_resourceRoot;
}

bool ResourceManager::BootstrapProjectAssets() {
    if (m_resourceRoot.Empty()) {
        return false;
    }

    const bool bootstrapped = m_assetImportService.BootstrapProject();
    m_projectAssetIndex.RefreshFrom(m_assetImportService);
    return bootstrapped;
}

void ResourceManager::AddRef(ResourceGUID guid) {
    std::lock_guard lock(m_mutex);
    
    auto* it = m_refCounts.Find(guid);
    if (it == nullptr) {
        m_refCounts.Insert(guid, 1);
    } else {
        (*it)++;
    }
    
    if (!m_resourceCache.Contains(guid)) {
        ReloadResource(guid);
    }
}

void ResourceManager::Release(ResourceGUID guid) {
    std::lock_guard lock(m_mutex);
    
    auto* it = m_refCounts.Find(guid);
    if (it != nullptr) {
        (*it)--;
        
        if (*it == 0) {
            m_refCounts.Erase(guid);
            m_cache.OnZeroRefCount(guid);
        }
    }
}

Core::uint32 ResourceManager::GetRefCount(ResourceGUID guid) const {
    auto* it = m_refCounts.Find(guid);
    return it != nullptr ? *it : 0;
}

void ResourceManager::RegisterLoader(IResourceLoader* loader) {
    std::lock_guard lock(m_mutex);
    m_loaders.Insert(loader->GetResourceType(), loader);
}

IResourceLoader* ResourceManager::GetLoader(ResourceType type) const {
    auto* it = m_loaders.Find(type);
    return it != nullptr ? *it : nullptr;
}

IResourceLoader* ResourceManager::FindLoader(ResourceType type) {
    return GetLoader(type);
}

IResource* ResourceManager::FindInCache(ResourceGUID guid) {
    std::lock_guard lock(m_mutex);
    
    auto* it = m_resourceCache.Find(guid);
    return it != nullptr ? *it : nullptr;
}

void ResourceManager::AddToCache(ResourceGUID guid, IResource* resource) {
    std::lock_guard lock(m_mutex);

    resource->m_guid = guid;
    m_resourceCache.Insert(guid, resource);
    m_memoryUsage += resource->GetMemorySize();
    m_cache.Add(guid, resource);
    
    if (m_memoryUsage > m_memoryBudget) {
        m_cache.OnMemoryPressure(m_memoryUsage - m_memoryBudget);
    }
}

void ResourceManager::Unload(ResourceGUID guid) {
    IResource* resource = nullptr;
    {
        std::lock_guard lock(m_mutex);

        auto* it = m_resourceCache.Find(guid);
        if (it != nullptr) {
            resource = *it;
            m_resourceCache.Erase(guid);
            m_cache.Remove(guid);
            m_guidToPath.Erase(guid);
            m_memoryUsage -= resource->GetMemorySize();
        }
    }

    if (resource != nullptr) {
        resource->Release();
    }
}

void ResourceManager::UnloadAll() {
    Containers::Array<IResource*> resourcesToRelease;
    {
        std::lock_guard lock(m_mutex);

        const auto cachedResources = m_resourceCache.GetPairs();
        resourcesToRelease.Reserve(cachedResources.Size());
        for (const auto& pair : cachedResources) {
            if (pair.second != nullptr) {
                resourcesToRelease.PushBack(pair.second);
            }
        }

        m_resourceCache.Clear();
        m_cache.Clear();
        m_refCounts.Clear();
        m_guidToPath.Clear();
        m_memoryUsage = 0;
    }

    for (IResource* resource : resourcesToRelease) {
        if (resource != nullptr) {
            resource->Release();
        }
    }
}

void ResourceManager::SetMemoryBudget(size_t bytes) {
    m_memoryBudget = bytes;
}

size_t ResourceManager::GetMemoryUsage() const {
    return m_memoryUsage;
}

size_t ResourceManager::GetMemoryBudget() const {
    return m_memoryBudget;
}

void ResourceManager::FlushCache() {
    m_cache.Flush();
}

IResource* ResourceManager::Find(const Containers::String& path) {
    return Find(ResourceGUID::Generate(path));
}

IResource* ResourceManager::Find(ResourceGUID guid) {
    return FindInCache(guid);
}

bool ResourceManager::Exists(const Containers::String& path) const {
    return Exists(ResourceGUID::Generate(path));
}

bool ResourceManager::Exists(ResourceGUID guid) const {
    return m_resourceCache.Contains(guid);
}

Containers::String ResourceManager::ResolvePath(const Containers::String& relativePath) const {
    return m_resourceRoot + "/" + relativePath;
}

void ResourceManager::LoadAsync(const Containers::String& path, ResourceType type,
                               std::function<void(LoadResult)> callback) {
    LoadAsync(path, type, nullptr, callback);
}

void ResourceManager::LoadAsync(const Containers::String& path, ResourceType type, 
                               ImportSettings* settings,
                               std::function<void(LoadResult)> callback) {
    EnsureInitialized();

    if (!m_asyncLoader) {
        if (callback) {
            callback(LoadResult("Async loader is not initialized"));
        }
        return;
    }

    m_asyncLoader->Submit(path, type, settings, std::move(callback));
}

void ResourceManager::UpdateAsyncLoads() {
    if (m_asyncLoader) {
        m_asyncLoader->Update();
    }
}

bool ResourceManager::IsAsyncLoading() const {
    return m_asyncLoader && m_asyncLoader->IsLoading();
}

Core::uint32 ResourceManager::GetAsyncPendingCount() const {
    return m_asyncLoader ? m_asyncLoader->GetPendingCount() : 0;
}

void ResourceManager::Unload(const Containers::String& path) {
    Unload(ResourceGUID::Generate(path));
}

void ResourceManager::UnloadUnused() {
    
}

void ResourceManager::UnregisterLoader(ResourceType type) {
    m_loaders.Erase(type);
}

void ResourceManager::ReloadResource(ResourceGUID guid) {
    auto* pathIt = m_guidToPath.Find(guid);
    if (pathIt == nullptr) {
        return;
    }

    const Containers::String path = *pathIt;
    auto* typeIt = m_resourceCache.Find(guid);
    (void)typeIt;
}

Containers::Array<Containers::String> ResourceManager::GetResourcePaths() const {
    Containers::Array<Containers::String> paths;
    for (const auto& pair : m_guidToPath) {
        paths.PushBack(pair.second);
    }
    return paths;
}

void ResourceManager::UnloadGroup(const Containers::Array<ResourceGUID>& guids) {
    Containers::Array<IResource*> resourcesToRelease;
    {
        std::lock_guard lock(m_mutex);
        resourcesToRelease.Reserve(guids.Size());
        for (const auto& guid : guids) {
            auto* it = m_resourceCache.Find(guid);
            if (it != nullptr) {
                IResource* resource = *it;
                m_resourceCache.Erase(guid);
                m_cache.Remove(guid);
                m_guidToPath.Erase(guid);
                m_memoryUsage -= resource->GetMemorySize();
                resourcesToRelease.PushBack(resource);
            }
        }
    }

    for (IResource* resource : resourcesToRelease) {
        if (resource != nullptr) {
            resource->Release();
        }
    }
}

void ResourceManager::RefreshProjectAssets() {
    if (!m_resourceRoot.Empty()) {
        m_assetImportService.Refresh();
        m_projectAssetIndex.RefreshFrom(m_assetImportService);
    }
}

bool ResourceManager::CanReimportProjectAsset(const Containers::String& path) const {
    if (m_resourceRoot.Empty() || path.Empty()) {
        return false;
    }

    ResourceType importType = ResourceType::Unknown;
    return m_assetImportService.TryGetImportableResourceType(path, importType);
}

bool ResourceManager::ReimportProjectAsset(const Containers::String& path) {
    if (m_resourceRoot.Empty() || path.Empty()) {
        return false;
    }

    UnloadAll();

    AssetImportService::ImportedAsset importedAsset;
    const bool reimported = m_assetImportService.ReimportAsset(path, importedAsset);
    m_projectAssetIndex.RefreshFrom(m_assetImportService);
    if (reimported && importedAsset.assetGuid.IsValid() && !importedAsset.relativePath.Empty()) {
        m_projectAssetIndex.RememberResolvedPath(importedAsset.assetGuid, importedAsset.relativePath);
    }

    return reimported;
}

bool ResourceManager::ClearProjectLibraryCache() {
    if (m_resourceRoot.Empty()) {
        return false;
    }

    UnloadAll();
    const bool cleared = m_assetImportService.ClearLibraryCache();
    m_projectAssetIndex.RefreshFrom(m_assetImportService);
    return cleared;
}

bool ResourceManager::RebuildProjectAssetCache() {
    if (m_resourceRoot.Empty()) {
        return false;
    }

    UnloadAll();
    const bool rebuilt = m_assetImportService.RebuildLibraryCache();
    m_projectAssetIndex.RefreshFrom(m_assetImportService);
    return rebuilt;
}

Containers::String ResourceManager::GetProjectLibraryRoot() const {
    return m_assetImportService.GetLibraryRoot();
}

AssetImportService::ImportStatusSnapshot ResourceManager::GetProjectAssetImportStatus() const {
    return m_assetImportService.GetLastImportStatus();
}

bool ResourceManager::TryGetAssetRef(const Containers::String& path, ResourceType resourceType, AssetRef& outRef) const {
    const bool resolved = m_projectAssetIndex.TryGetAssetRef(m_assetImportService, path, resourceType, outRef);

    if (ShouldTraceResourcePath(path)) {
        Debug::Logger::Get().Info(
            Debug::LogCategory::FileSystem,
            Containers::String("[ResourceManager] TryGetAssetRef path=") +
            path +
            " type=" +
            GetResourceTypeName(resourceType) +
            " success=" +
            Containers::String(resolved ? "1" : "0") +
            " ref=" +
            EncodeAssetRef(outRef));
    }
    return resolved;
}

bool ResourceManager::TryResolveAssetPath(const AssetRef& assetRef, Containers::String& outPath) const {
    const bool resolved = m_projectAssetIndex.TryResolveAssetPath(m_assetImportService, assetRef, outPath);

    if (resolved && ShouldTraceResourcePath(outPath)) {
        Debug::Logger::Get().Info(
            Debug::LogCategory::FileSystem,
            Containers::String("[ResourceManager] TryResolveAssetPath ref=") +
            EncodeAssetRef(assetRef) +
            " path=" +
            outPath);
    }
    return resolved;
}

void ResourceManager::BeginDeferredSceneLoad() {
    ++m_deferredSceneLoadDepth;
}

void ResourceManager::EndDeferredSceneLoad() {
    const Core::uint32 currentDepth = m_deferredSceneLoadDepth.load();
    if (currentDepth == 0) {
        return;
    }

    --m_deferredSceneLoadDepth;
}

bool ResourceManager::IsDeferredSceneLoadEnabled() const {
    return m_deferredSceneLoadDepth.load() > 0;
}

LoadResult ResourceManager::LoadResource(const Containers::String& path,
                                         ResourceType type,
                                         ImportSettings* settings) {
    EnsureInitialized();

    const ResourceGUID guid = ResourceGUID::Generate(path);

    if (ShouldTraceResourcePath(path)) {
        Debug::Logger::Get().Info(
            Debug::LogCategory::FileSystem,
            Containers::String("[ResourceManager] LoadResource request path=") +
            path +
            " type=" +
            GetResourceTypeName(type) +
            " root=" +
            m_resourceRoot);
    }

    if (IResource* cached = FindInCache(guid)) {
        if (ShouldTraceResourcePath(path)) {
            Debug::Logger::Get().Info(
                Debug::LogCategory::FileSystem,
                Containers::String("[ResourceManager] LoadResource cache-hit path=") + path);
        }
        return LoadResult(cached);
    }

    IResourceLoader* loader = FindLoader(type);
    if (loader == nullptr) {
        Debug::Logger::Get().Warning(Debug::LogCategory::FileSystem,
            Containers::String("No loader found for resource type: ") +
            GetResourceTypeName(type));
        return LoadResult(false, "Loader not found");
    }

    const InFlightLoadKey inFlightKey{ guid, type };
    std::shared_ptr<InFlightLoadState> inFlightState;
    bool shouldExecuteLoad = false;
    {
        std::unique_lock<std::mutex> inFlightLock(m_inFlightLoadsMutex);
        auto inFlightIt = m_inFlightLoads.find(inFlightKey);
        if (inFlightIt == m_inFlightLoads.end()) {
            inFlightState = std::make_shared<InFlightLoadState>();
            m_inFlightLoads.emplace(inFlightKey, inFlightState);
            shouldExecuteLoad = true;
        } else {
            inFlightState = inFlightIt->second;
        }
    }

    auto completeInFlightLoad = [&](const LoadResult& result) {
        {
            std::lock_guard<std::mutex> inFlightLock(m_inFlightLoadsMutex);
            inFlightState->completed = true;
            inFlightState->success = result && result.resource != nullptr;
            inFlightState->errorMessage = result.errorMessage;
            m_inFlightLoads.erase(inFlightKey);
        }
        inFlightState->condition.notify_all();
    };

    if (!shouldExecuteLoad) {
        if (ShouldTraceResourcePath(path)) {
            Debug::Logger::Get().Info(
                Debug::LogCategory::FileSystem,
                Containers::String("[ResourceManager] LoadResource wait-inflight path=") +
                path +
                " type=" +
                GetResourceTypeName(type));
        }

        {
            std::unique_lock<std::mutex> inFlightLock(m_inFlightLoadsMutex);
            inFlightState->condition.wait(inFlightLock, [&inFlightState]() {
                return inFlightState->completed;
            });
        }

        if (IResource* cached = FindInCache(guid)) {
            if (ShouldTraceResourcePath(path)) {
                Debug::Logger::Get().Info(
                    Debug::LogCategory::FileSystem,
                    Containers::String("[ResourceManager] LoadResource in-flight-cache-hit path=") + path);
            }
            return LoadResult(cached);
        }

        return LoadResult(
            !inFlightState->errorMessage.Empty()
                ? inFlightState->errorMessage
                : Containers::String("In-flight load completed without cached resource"));
    }

    Containers::String loadPath = path;
    AssetImportService::ImportedAsset resolvedAsset;
    ResourceType importableType = ResourceType::Unknown;
    const bool shouldUseProjectArtifact =
        !m_resourceRoot.Empty() &&
        m_assetImportService.TryGetImportableResourceType(path, importableType) &&
        importableType == type;

    if (shouldUseProjectArtifact &&
        m_assetImportService.EnsureArtifact(path, type, resolvedAsset) &&
        resolvedAsset.artifactReady) {
        m_projectAssetIndex.RememberResolvedPath(resolvedAsset.assetGuid, resolvedAsset.relativePath);
        loadPath = resolvedAsset.runtimeLoadPath;
        if (ShouldTraceResourcePath(path)) {
            Debug::Logger::Get().Info(
                Debug::LogCategory::FileSystem,
                Containers::String("[ResourceManager] LoadResource artifact path=") +
                path +
                " artifact=" +
                loadPath);
        }
    } else if (ShouldTraceResourcePath(path)) {
        Debug::Logger::Get().Info(
            Debug::LogCategory::FileSystem,
            Containers::String("[ResourceManager] LoadResource direct path=") +
            path +
            " loadPath=" +
            loadPath);
    }

    LoadResult result;
    try {
        result = loader->Load(loadPath, settings);
    } catch (const std::exception& exception) {
        result = LoadResult(
            Containers::String("LoadResource exception: ") +
            Containers::String(exception.what()));
    } catch (...) {
        result = LoadResult("LoadResource exception: unknown");
    }

    if (!result || result.resource == nullptr) {
        Debug::Logger::Get().Error(Debug::LogCategory::FileSystem,
            Containers::String("Failed to load resource: ") + path + " - " + result.errorMessage);
        completeInFlightLoad(result);
        return result;
    }

    if (ShouldTraceResourcePath(path)) {
        Debug::Logger::Get().Info(
            Debug::LogCategory::FileSystem,
            Containers::String("[ResourceManager] LoadResource success path=") +
            path +
            " loadPath=" +
            loadPath);
    }

    result.resource->m_path = path;
    AddToCache(guid, result.resource);
    {
        std::lock_guard lock(m_mutex);
        m_guidToPath.Insert(guid, path);
    }
    completeInFlightLoad(result);
    return result;
}

} // namespace Resources
} // namespace XCEngine