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>

namespace XCEngine {
namespace Resources {

namespace {

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;

} // namespace

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

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

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

    RegisterBuiltinLoader(*this, g_materialLoader);
    RegisterBuiltinLoader(*this, g_meshLoader);
    RegisterBuiltinLoader(*this, g_shaderLoader);
    RegisterBuiltinLoader(*this, g_textureLoader);
}

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

void ResourceManager::SetResourceRoot(const Containers::String& rootPath) {
    m_resourceRoot = rootPath;
    if (!m_resourceRoot.Empty()) {
        ResourceFileSystem::Get().Initialize(rootPath);
        m_assetDatabase.Initialize(rootPath);
    } else {
        ResourceFileSystem::Get().Shutdown();
        m_assetDatabase.Shutdown();
    }
}

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

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) {
    std::lock_guard lock(m_mutex);
    
    auto* it = m_resourceCache.Find(guid);
    if (it != nullptr) {
        IResource* resource = *it;
        m_resourceCache.Erase(guid);
        m_guidToPath.Erase(guid);
        m_memoryUsage -= resource->GetMemorySize();
        resource->Release();
    }
}

void ResourceManager::UnloadAll() {
    std::lock_guard lock(m_mutex);

    const auto cachedResources = m_resourceCache.GetPairs();
    for (const auto& pair : cachedResources) {
        if (pair.second != nullptr) {
            pair.second->Release();
        }
    }
    m_resourceCache.Clear();
    m_refCounts.Clear();
    m_guidToPath.Clear();
    m_memoryUsage = 0;
}

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) {
    m_asyncLoader->Submit(path, type, settings, callback);
}

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) {
    std::lock_guard lock(m_mutex);
    for (const auto& guid : guids) {
        auto* it = m_resourceCache.Find(guid);
        if (it != nullptr) {
            IResource* resource = *it;
            m_resourceCache.Erase(guid);
            m_guidToPath.Erase(guid);
            m_memoryUsage -= resource->GetMemorySize();
            resource->Release();
        }
    }
}

void ResourceManager::RefreshAssetDatabase() {
    if (!m_resourceRoot.Empty()) {
        m_assetDatabase.Refresh();
    }
}

bool ResourceManager::TryGetAssetRef(const Containers::String& path, ResourceType resourceType, AssetRef& outRef) const {
    return m_assetDatabase.TryGetAssetRef(path, resourceType, outRef);
}

LoadResult ResourceManager::LoadResource(const Containers::String& path,
                                         ResourceType type,
                                         ImportSettings* settings) {
    const ResourceGUID guid = ResourceGUID::Generate(path);

    if (IResource* cached = FindInCache(guid)) {
        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");
    }

    Containers::String loadPath = path;
    AssetDatabase::ResolvedAsset resolvedAsset;
    if (!m_resourceRoot.Empty() &&
        m_assetDatabase.EnsureArtifact(path, type, resolvedAsset) &&
        resolvedAsset.artifactReady) {
        loadPath = resolvedAsset.artifactMainPath;
    }

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

    result.resource->m_path = path;
    AddToCache(guid, result.resource);
    m_guidToPath.Insert(guid, path);
    return result;
}

} // namespace Resources
} // namespace XCEngine