LuaCsForBarotraumaEP/Barotrauma/BarotraumaClient/ClientSource/Steam/Workshop.cs

#nullable enable
using Barotrauma.IO;
using Microsoft.Xna.Framework.Graphics;
using RestSharp;
using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Barotrauma.Extensions;

namespace Barotrauma.Steam
{
    static partial class SteamManager
    {
        public static partial class Workshop
        {
            public const int MaxThumbnailSize = 1024 * 1024;

            /// <summary>
            /// Tags the players can choose for their workshop items. These must match the ones defined in the Steamworks backend. They're case insensitive, but must otherwise match exactly for the tag filtering to work correctly.
            /// The localized names for these are fetched from the loca files with the identifier "workshop.contenttag.{tag.RemoveWhitespace()}".
            /// </summary>
            public static readonly ImmutableArray<Identifier> Tags = new []
            {
                "submarine",
                "item",
                "monster",
                "mission",
                "outpost",
                "beacon station",
                "wreck",
                "ruin",
                "weapons",
                "medical",
                "equipment",
                "art",
                "event set",
                "total conversion",
                "game mode",
                "gameplay mechanics",
                "environment",
                "item assembly",
                "language",
                "qol",
                "client-side",
                "server-side",
                "outdated",
                "library"
            }.ToIdentifiers().ToImmutableArray();
            
            public class ItemThumbnail : IDisposable
            {
                private struct RefCounter
                {
                    internal bool Loading;
                    internal Texture2D? Texture;
                    internal int Count;
                }
                private readonly static Dictionary<UInt64, RefCounter> TextureRefs
                    = new Dictionary<ulong, RefCounter>();

                public UInt64 ItemId { get; private set; }
                public Texture2D? Texture
                {
                    get
                    {
                        lock (TextureRefs)
                        {
                            if (TextureRefs.TryGetValue(ItemId, out var refCounter))
                            {
                                return refCounter.Texture;
                            }
                        }
                        return null;
                    }
                }

                public bool Loading
                {
                    get
                    {
                        lock (TextureRefs)
                        {
                            if (TextureRefs.TryGetValue(ItemId, out var refCounter))
                            {
                                return refCounter.Loading;
                            }
                        }
                        return false;
                    }
                }

                public ItemThumbnail(in Steamworks.Ugc.Item item, CancellationToken cancellationToken)
                {
                    ItemId = item.Id;
                    lock (TextureRefs)
                    {
                        if (TextureRefs.TryGetValue(ItemId, out var refCounter))
                        {
                            TextureRefs[ItemId] = new RefCounter { Texture = refCounter.Texture, Count = refCounter.Count + 1, Loading = refCounter.Loading };
                        }
                        else
                        {
                            TextureRefs[ItemId] = new RefCounter { Texture = null, Count = 1, Loading = true };
                            TaskPool.Add($"Workshop thumbnail {item.Title}", GetTexture(item, cancellationToken), SaveTextureToRefCounter(item.Id));
                        }
                    }
                }

                public void Dispose()
                {
                    if (ItemId == 0) { return; }
                    lock (TextureRefs)
                    {
                        var refCounter = TextureRefs[ItemId];
                        TextureRefs[ItemId] = new RefCounter { Texture = refCounter.Texture, Count = refCounter.Count - 1 };
                        if (TextureRefs[ItemId].Count <= 0)
                        {
                            TextureRefs[ItemId].Texture?.Dispose();
                            TextureRefs.Remove(ItemId);
                        }
                        ItemId = 0;
                    }
                }

                private static async Task<Texture2D?> GetTexture(Steamworks.Ugc.Item item, CancellationToken cancellationToken)
                {
                    await Task.Yield();

                    string? thumbnailUrl = item.PreviewImageUrl;
                    if (thumbnailUrl.IsNullOrWhiteSpace()) { return null; }
                    var client = RestFactory.CreateClient(thumbnailUrl);
                    var request = RestFactory.CreateRequest(".");
                    IRestResponse response = await client.ExecuteAsync(request, cancellationToken);
                    if (response.ErrorException != null)
                    {
                        DebugConsole.NewMessage($"Connection error: Failed to load workshop item thumbnail for {item.Id} ({response.ErrorException.Message}).");
                    }
                    else if (response is { StatusCode: System.Net.HttpStatusCode.OK, ResponseStatus: ResponseStatus.Completed })
                    {
                        using var dataStream = new System.IO.MemoryStream();
                        await dataStream.WriteAsync(response.RawBytes, cancellationToken);
                        dataStream.Seek(0, System.IO.SeekOrigin.Begin);
                        return TextureLoader.FromStream(dataStream, compress: false);
                    }
                    return null;
                }

                private static Action<Task> SaveTextureToRefCounter(UInt64 itemId)
                    => (t) =>
                    {
                        if (t.IsCanceled) { return; }
                        Texture2D? texture = ((Task<Texture2D?>)t).Result;
                        lock (TextureRefs)
                        {
                            if (TextureRefs.TryGetValue(itemId, out var refCounter))
                            {
                                TextureRefs[itemId] = new RefCounter { Texture = texture, Count = refCounter.Count, Loading = false };
                            }
                            else if (texture != null)
                            {
                                texture.Dispose();
                            }
                        }
                    };

                public override int GetHashCode() => (int)ItemId;

                public override bool Equals(object? obj)
                    => obj is ItemThumbnail { ItemId: UInt64 otherId }
                        && otherId == ItemId;
            }

            public const string PublishStagingDir = "WorkshopStaging";

            public static void DeletePublishStagingCopy()
            {
                if (Directory.Exists(PublishStagingDir)) { Directory.Delete(PublishStagingDir, recursive: true); }
            }

            private static void RefreshLocalMods()
            {
                CrossThread.RequestExecutionOnMainThread(() => ContentPackageManager.LocalPackages.Refresh());
            }
            
            public static async Task CreatePublishStagingCopy(string title, string modVersion, ContentPackage contentPackage)
            {
                await Task.Yield();
                
                if (!ContentPackageManager.LocalPackages.Contains(contentPackage))
                {
                    throw new Exception("Expected local package");
                }

                DeletePublishStagingCopy();
                Directory.CreateDirectory(PublishStagingDir, catchUnauthorizedAccessExceptions: false);
                await CopyDirectory(contentPackage.Dir, contentPackage.Name, Path.GetDirectoryName(contentPackage.Path)!, PublishStagingDir, ShouldCorrectPaths.No);

                var stagingFileListPath = Path.Combine(PublishStagingDir, ContentPackage.FileListFileName);

                var result = ContentPackage.TryLoad(stagingFileListPath);
                if (!result.TryUnwrapSuccess(out var tempPkg))
                {
                    throw new Exception("Staging copy could not be loaded",
                        result.TryUnwrapFailure(out var exception) ? exception : null);
                }

                //Load filelist.xml and write the hash into it so anyone downloading this mod knows what it should be                
                ModProject modProject = new ModProject(tempPkg)
                {
                    ModVersion = modVersion,
                    Name = title,
                    ExpectedHash = tempPkg.CalculateHash(name: title, modVersion: modVersion)
                };
                modProject.Save(stagingFileListPath);
            }

            public static async Task<Option<ContentPackage>> CreateLocalCopy(ContentPackage contentPackage)
            {
                await Task.Yield();
                
                if (!ContentPackageManager.WorkshopPackages.Contains(contentPackage))
                {
                    throw new Exception("Expected Workshop package");
                }

                if (!contentPackage.UgcId.TryUnwrap(out var ugcId) || !(ugcId is SteamWorkshopId workshopId))
                {
                    throw new Exception($"Steam Workshop ID not set for {contentPackage.Name}");
                }

                string sanitizedName = ToolBox.RemoveInvalidFileNameChars(contentPackage.Name).Trim();
                if (sanitizedName.IsNullOrWhiteSpace())
                {
                    throw new Exception($"Sanitized name for {contentPackage.Name} is empty");
                }

                string newPath = $"{ContentPackage.LocalModsDir}/{sanitizedName}";
                if (File.Exists(newPath) || Directory.Exists(newPath))
                {
                    newPath += $"_{workshopId.Value}";
                }

                if (File.Exists(newPath) || Directory.Exists(newPath))
                {
                    throw new Exception($"{newPath} already exists");
                }

                await CopyDirectory(contentPackage.Dir, contentPackage.Name, Path.GetDirectoryName(contentPackage.Path)!, newPath, ShouldCorrectPaths.Yes);

                ModProject modProject = new ModProject(contentPackage);
                modProject.DiscardHashAndInstallTime();
                modProject.Save(Path.Combine(newPath, ContentPackage.FileListFileName));

                RefreshLocalMods();

                return ContentPackageManager.LocalPackages.FirstOrNone(p => p.UgcId == contentPackage.UgcId);
            }

            private struct InstallWaiter
            {
                private static readonly HashSet<ulong> waitingIds = new HashSet<ulong>();
                public ulong Id { get; private set; }

                public InstallWaiter(ulong id)
                {
                    Id = id;
                    lock (waitingIds) { waitingIds.Add(Id); }
                }

                public bool Waiting
                {
                    get
                    {
                        if (Id == 0) { return false; }

                        lock (waitingIds)
                        {
                            return waitingIds.Contains(Id);
                        }
                    }
                }

                public static void StopWaiting(ulong id)
                {
                    lock (waitingIds)
                    {
                        waitingIds.Remove(id);
                    }
                }
            }

            public static async Task Reinstall(Steamworks.Ugc.Item workshopItem)
            {
                NukeDownload(workshopItem);
                var toUninstall
                    = ContentPackageManager.WorkshopPackages.Where(p =>
                            p.UgcId.TryUnwrap(out var ugcId)
                            && ugcId is SteamWorkshopId workshopId
                            && workshopId.Value == workshopItem.Id)
                        .ToHashSet();
                toUninstall.Select(p => p.Dir).ForEach(d => Directory.Delete(d));
                CrossThread.RequestExecutionOnMainThread(() => ContentPackageManager.WorkshopPackages.Refresh());
                var installWaiter = WaitForInstall(workshopItem);
                DownloadModThenEnqueueInstall(workshopItem);
                await installWaiter;
            }

            public static async Task WaitForInstall(Steamworks.Ugc.Item item)
                => await WaitForInstall(item.Id);
            
            public static async Task WaitForInstall(ulong item)
            {
                var installWaiter = new InstallWaiter(item);
                while (installWaiter.Waiting) { await Task.Delay(500); }
                await Task.Delay(500);
            }
            
            public static void OnItemDownloadComplete(ulong id, bool forceInstall = false)
            {
                if (Screen.Selected is not MainMenuScreen && !forceInstall)
                {
                    if (!MainMenuScreen.WorkshopItemsToUpdate.Contains(id))
                    {
                        MainMenuScreen.WorkshopItemsToUpdate.Enqueue(id);
                    }
                    return;
                }
                else if (!CanBeInstalled(id))
                {
                    DebugConsole.Log($"Cannot install {id}");
                    InstallWaiter.StopWaiting(id);
                }
                else if (ContentPackageManager.WorkshopPackages.Any(p =>
                        p.UgcId.TryUnwrap(out var ugcId)
                        && ugcId is SteamWorkshopId workshopId
                        && workshopId.Value == id))
                {
                    DebugConsole.Log($"Already installed {id}.");
                    InstallWaiter.StopWaiting(id);
                }
                else if (InstallTaskCounter.IsInstalling(id))
                {
                    DebugConsole.Log($"Already installing {id}.");
                }
                else
                {
                    DebugConsole.Log($"Finished downloading {id}, installing...");
                    TaskPool.Add($"InstallItem{id}", InstallMod(id), t => InstallWaiter.StopWaiting(id));
                }
            }
        }
    }
}