AssetService.cs

Go to the documentation of this file.

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Reflection;
using System.Threading;
using System.Threading.Tasks;
using BepInEx;
using BepInEx.Bootstrap;
using UnityEngine;
using UnityEngine.UIElements;
using WKLib.API;
using WKLib.Utilities;
using Object = UnityEngine.Object;
 
namespace WKLib.API.Assets;
 
public class AssetService
{
    private readonly WKLibAPI _wkLibAPI;
    public readonly string AssemblyFolder;
    private readonly Dictionary<string, AssetBundle> _bundleCache = new();
    private readonly Dictionary<string, Dictionary<string, M_Level>> _loadedLevelsCache = new();
    private readonly Dictionary<string, M_Gamemode> _gamemodeCache = new();
    
    private readonly CancellationTokenSource _sQuitCts = new();
    
    public AssetService(WKLibAPI API)
    {
        _wkLibAPI = API ?? throw new ArgumentNullException(nameof(API));
        if (!Chainloader.PluginInfos.TryGetValue(API.GUID, out var pluginInfo))
            throw new Exception($"WKLibAPI GUID Mistmatch with PluginInfo, {API.DisplayName}, {API.GUID}");
        
        var assemblyPath = pluginInfo.Location;
        AssemblyFolder = Path.GetDirectoryName(assemblyPath) ?? string.Empty;
 
        Application.quitting += () =>
        {
            WKLog.Debug("[AssetService] Quitting...");
            _sQuitCts.Cancel();
        };
    }
    
    #region Loading Asset Bundles
    
    public async Task<AssetBundle> LoadBundleRelativeAsync(
        string relativePath,
        IProgress<float> progress = null)
    {
        progress?.Report(0f);
        if (string.IsNullOrEmpty(relativePath))
        {
            WKLog.Error("[AssetService] Path is null or empty.");
            progress?.Report(1f);
            return null;
        }
 
        var cacheKey = $"{_wkLibAPI.DisplayName}:{Path.GetFileNameWithoutExtension(relativePath)}";
        if (_bundleCache.TryGetValue(cacheKey, out var cached) && cached is not null)
        {
            WKLog.Debug($"[AssetService] Returning cached bundle: {cacheKey}");
            progress?.Report(1f);
            return cached;
        }
 
        var fullPath = Path.Combine(AssemblyFolder, relativePath);
        if (!File.Exists(fullPath))
        {
            WKLog.Error($"[AssetService] Bundle not found at {fullPath}");
            progress?.Report(1f);
            return null;
        }
 
        var request = AssetBundle.LoadFromFileAsync(fullPath);
        if (request is null)
        {
            WKLog.Error($"[AssetService] Failed to load bundle: {fullPath} (request)");
            progress?.Report(1f);
            return null;
        }
 
        try
        {
            while (!request.isDone)
            {
                _sQuitCts.Token.ThrowIfCancellationRequested();
 
                progress?.Report(request.progress * 1.1f);
                await Task.Yield();
            }
        }
        catch (OperationCanceledException)
        {
            WKLog.Debug($"[AssetService] Asset Bundle Load cancelled");
            return null;
        }
        
        var bundle = request.assetBundle;
        if (bundle is null)
        {
            WKLog.Error($"[AssetService] Failed to load bundle at {fullPath} (bundle)");
            progress?.Report(1f);
            return null;
        }
 
        _bundleCache[cacheKey] = bundle;
        WKLog.Debug($"[AssetService] Loaded & cached bundle: {cacheKey}");
        progress?.Report(1f);
        return bundle;
    }
    
    #endregion
    
    #region Unloading Asset Bundles
    
    public void UnloadBundleRelative(string relativePath, bool unloadAllLoadedObjects = false)
    {
        var cacheKey = $"{_wkLibAPI.DisplayName}:{relativePath}";
        if (_bundleCache.TryGetValue(cacheKey, out var bundle) && bundle is not null)
        {
            try
            {
                bundle.Unload(unloadAllLoadedObjects);
            }
            catch
            {
                WKLog.Error($"[AssetService] Failed to unload bundle: {cacheKey}");
                return;
            }
            WKLog.Debug($"[AssetService] Unloaded bundle: {cacheKey}");
            _bundleCache.Remove(cacheKey);
        }
    }
    
    public void UnloadAllBundles(bool unloadAllLoadedObjects = false)
    {
        var keys = _bundleCache.Keys.Where(k => k.StartsWith(_wkLibAPI.DisplayName + ":")).ToList();
        foreach (var key in keys)
        {
            if (_bundleCache[key] is not null)
            {
                try
                {
                    _bundleCache[key].Unload(unloadAllLoadedObjects);
                }
                catch
                {
                    WKLog.Error($"[AssetService] Failed to unload bundle: {key}");
                    continue;
                }
                WKLog.Debug($"[AssetService] Unloaded bundle: {key}");
            }
            _bundleCache.Remove(key);
        }
    }
    
    #endregion
    
    #region Loading Levels
    
    public async Task<Dictionary<string, M_Level>> LoadAllLevelsFromBundle(
        AssetBundle bundle,
        IProgress<float> progress = null)
    {
        var foundLevels = new Dictionary<string, M_Level>();
        
        if (bundle is null)
        {
            WKLog.Error("[AssetService] LoadAllLevelsFromBundle called with null bundle.");
            progress?.Report(1f);
            return foundLevels;
        }
 
        var cacheKey = $"{_wkLibAPI.DisplayName}:{bundle.name}";
 
        if (_loadedLevelsCache.ContainsKey(cacheKey))
        {
            WKLog.Debug($"[AssetService] Returning Cached loaded levels for bundle: {cacheKey}");
            progress?.Report(1f);
            return _loadedLevelsCache.GetValueOrDefault(cacheKey);
        }
 
        var request = bundle.LoadAllAssetsAsync<GameObject>();
        try
        {
            while (!request.isDone)
            {
                _sQuitCts.Token.ThrowIfCancellationRequested();
 
                progress?.Report(request.progress * 0.5f);
                await Task.Yield();
            }
        }
        catch (OperationCanceledException)
        {
            WKLog.Debug($"[AssetService] LoadAllLevelsFromBundle canceled.");
            return null;
        }
        
 
        var currentDatabase = CL_AssetManager.GetDatabase(_wkLibAPI.DisplayName);
        if (currentDatabase == null)
        {
            WKLog.Info($"[AssetService] Creating new asset database for {_wkLibAPI.DisplayName}");
            currentDatabase = new WKAssetDatabase();
 
            CL_AssetManager.AddNewDatabase(_wkLibAPI.DisplayName, currentDatabase);
        }
 
        var allGOs = request.allAssets.Cast<GameObject>().ToList();
        var total = allGOs.Count;
        const int batchSize = 20;
 
        for (var i = 0; i < total; i += batchSize)
        {
            for (var j = i; j < i + batchSize && j < total; j++)
            {
                var go = allGOs[j]; 
                if (go.TryGetComponent<M_Level>(out var level))
                {
                    bool alreadyExists = false;
                    foreach (var levelAsset in currentDatabase.levelAssets)
                    {
                        if (levelAsset == null)
                            continue;
                        
                        if (levelAsset.id == level.levelName)
                        {
                            alreadyExists = true;
                            break;
                        }
                    }
                    
                    if (!alreadyExists)
                        currentDatabase.levelAssets.Add(M_Level.LevelAssetHolder.GetNewHolderFromLevel(level));
                    
                    // Log its name for debugging
                    WKLog.Debug($"[AssetService] Found level: {level.name}");
                    
                    if (!foundLevels.TryAdd(level.name, level))
                        WKLog.Debug($"[AssetService] Duplicate prefab.name '{level.name}' found; skipping the duplicate.");
                }
            }
            progress?.Report(0.5f + (i / (float)total) * 0.5f);
            await Task.Yield();
        }
        
        _loadedLevelsCache[cacheKey] = foundLevels;
        progress?.Report(1f);
        return foundLevels;
    }
    
    public List<M_Level> FindLevelsByName(string nameContains)
    {
        return CL_AssetManager.GetFullCombinedAssetDatabase().levelAssets
            .Select(levelAsset => levelAsset.level)
            .Where(level => level is not null)
            .Where(level => level.name.Contains(nameContains))
            .ToList();
    }
    
    #endregion
    
    #region Loading Gamemodes
    
    public async Task<M_Gamemode> LoadGameModeFromBundle(
        AssetBundle bundle,
        string assetName,
        IProgress<float> progress = null)
    {
        if (bundle is null)
        {
            WKLog.Error("[AssetService] Bundle is null; cannot load gamemode.");
            progress?.Report(1f);
            return null;
        }
 
        var cacheKey = $"{_wkLibAPI.DisplayName}:{assetName}";
        
        if (_gamemodeCache.TryGetValue(cacheKey, out var cachedGamemode))
        {
            WKLog.Debug($"[AssetService] Returning cached gamemode: '{cacheKey}'");
            progress?.Report(1f);
            return cachedGamemode;
        }
 
        var request = bundle.LoadAssetAsync<ScriptableObject>(assetName);
        while (!request.isDone)
        {
            progress?.Report(request.progress);
            await Task.Yield();
        }
        
        var gamemode = request.asset as M_Gamemode;
 
        if (gamemode is null)
        {
            WKLog.Error($"[AssetService] Gamemode '{assetName}' not found in bundle {bundle.name}");
            progress?.Report(1f);
            return null;
        }
        
        if (string.IsNullOrEmpty(gamemode.unlockAchievement))
            gamemode.unlockAchievement = "ACH_TUTORIAL";
        
        gamemode.gamemodePanel = Resources.FindObjectsOfTypeAll<UI_GamemodeScreen_Panel>()
            .FirstOrDefault(x => x.name == "Gamemode_Panel_Base");
        gamemode.loseScreen = Resources.FindObjectsOfTypeAll<UI_ScoreScreen>()
            .FirstOrDefault(x => x.name == "ScorePanel_Standard_Death");
        gamemode.winScreen = Resources.FindObjectsOfTypeAll<UI_ScoreScreen>()
            .FirstOrDefault(x => x.name == "ScorePanel_Standard_Win");
        
        _gamemodeCache[cacheKey] = gamemode;
        progress?.Report(1f);
        WKLog.Debug($"[AssetService] Loaded and cached gamemode: '{assetName}'");
        return gamemode;
    }
    
    #endregion
    
    #region Helper Methods
    
    public Sprite LoadPngAsSpriteRelative(string relativePngPath)
    {
        var fullPath = Path.Combine(AssemblyFolder, relativePngPath);
        if (!File.Exists(fullPath))
        {
            WKLog.Error($"[AssetService] PNG not found at: {fullPath}");
            return null;
        }
        var data = File.ReadAllBytes(fullPath);
        var tex = new Texture2D(2, 2, TextureFormat.RGBA32, false);
        if (!tex.LoadImage(data))
        {
            WKLog.Error($"[AssetService] Failed to load image: {relativePngPath}");
            return null;
        }
        var sprite = Sprite.Create(tex, new Rect(0,0,tex.width,tex.height), new Vector2(0.5f,0.5f));
        sprite.name = Path.GetFileNameWithoutExtension(relativePngPath);
        return sprite;
    }
    
    public Sprite LoadPngAsSprite(string pngFileName)
    {
        var tex = LoadPngTexture(pngFileName);
        if (tex is null) return null;
 
        var sprite = Sprite.Create(
            tex,
            new Rect(0, 0, tex.width, tex.height),
            new Vector2(0.5f, 0.5f)
        );
        sprite.name = pngFileName.Split(".png")[0];
        return sprite;
    }
    
    private Texture2D LoadPngTexture(string pngFileName)
    {
        var pngPath = Path.Combine(AssemblyFolder, "Assets", pngFileName);
        if (!File.Exists(pngPath))
        {
            WKLog.Error($"[AssetService] PNG not found at: {pngPath}");
            return null;
        }
 
        var data = File.ReadAllBytes(pngPath);
        var tex = new Texture2D(2, 2, TextureFormat.RGBA32, false);
        
        if (tex.LoadImage(data)) return tex;
        
        WKLog.Error($"[AssetService] Failed to decode PNG: {pngPath}");
        return null;
    }
    
    #endregion
}