AssetsFile.cs

using ICSharpCode.SharpZipLib.Zip;
using Microsoft.Extensions.Logging;
using Newtonsoft.Json;
 
namespace Gondwana.Assets;

[JsonObject(IsReference = true)]
public sealed class AssetsFile : IDisposable
{
    private static readonly List<AssetsFile> _allAssetsFiles = new();

    public static IReadOnlyList<AssetsFile> AllAssetsFiles => _allAssetsFiles.AsReadOnly();

    public static void ClearAll()
    {
        foreach (var assetFile in _allAssetsFiles.ToList())
            assetFile.Dispose();
    }
 
    private ZipFile? _zipFile;
 
    // In-memory source of truth.
    // Each entry stores its complete bytes independently of any live ZipFile handle.
    private readonly Dictionary<AssetsFileEntry, byte[]> _zipEntries = new();
 
    private bool _isLoaded;

    public static void Validate(string path, string? password = null, bool testData = true)
    {
        using var archive = new ZipFile(File.OpenRead(path));
        archive.Password = password;
        var keys = new HashSet<AssetsFileEntry>();
        foreach (ZipEntry entry in archive)
        {
            if (!entry.IsFile) continue;
            var key = AssetsFileEntry.FromString(entry.Name);
            if (key is null || !Enum.IsDefined(key.AssetType) || string.IsNullOrWhiteSpace(key.AssetName))
                throw new InvalidDataException($"Invalid bundle entry key: {entry.Name}");
            if (!keys.Add(key)) throw new InvalidDataException($"Duplicate bundle entry key: {entry.Name}");
        }
        if (!archive.TestArchive(testData)) throw new InvalidDataException("Bundle integrity check failed. Check the password and archive contents.");
    }
 
    [JsonConstructor]
    private AssetsFile() : this(register: true)
    {
    }
 
    private AssetsFile(bool register)
    {
        if (register) _allAssetsFiles.Add(this);
    }

    public static AssetsFile LoadOrCreate(string path, string? password = null, bool encrypt = false)
        => LoadOrCreate(path, password, encrypt, register: true);

    public static AssetsFile LoadOrCreate(string path, string? password, bool encrypt, bool register)
    {
        if (string.IsNullOrWhiteSpace(path))
            throw new ArgumentException("path cannot be null or empty.", nameof(path));
 
        var assetFile = new AssetsFile(register)
        {
            FilePath = path,
            Password = password,
            UseEncryption = encrypt
        };
 
        try
        {
            if (File.Exists(path))
                assetFile.LoadZip();
            else
                assetFile._isLoaded = true;
            return assetFile;
        }
        catch
        {
            assetFile.Dispose();
            throw;
        }
    }

    public static AssetsFile Load(Stream stream, string? password = null, bool register = true)
        => Load(stream, password, register, sourcePath: null);

    public static AssetsFile Load(Stream stream, string? password, bool register, string? sourcePath)
    {
        ArgumentNullException.ThrowIfNull(stream);
        if (!stream.CanRead)
            throw new ArgumentException("Stream must be readable.", nameof(stream));
 
        var assetFile = new AssetsFile(register)
        {
            Password = password,
            SourcePath = string.IsNullOrWhiteSpace(sourcePath) ? null : sourcePath
        };
 
        try
        {
            assetFile.LoadZip(stream);
            return assetFile;
        }
        catch
        {
            assetFile.Dispose();
            throw;
        }
    }

    [JsonProperty]
    public string FilePath { get; private set; } = string.Empty;

    [JsonProperty]
    public string? Password { get; private set; }

    [JsonProperty]
    public bool UseEncryption { get; private set; }

    [JsonIgnore]
    public string? SourcePath { get => string.IsNullOrWhiteSpace(FilePath) ? _sourcePath : FilePath; private set => _sourcePath = value; }
    private string? _sourcePath;
 
    private void EnsureLoaded()
    {
        if (!_isLoaded)
            LoadZip();
    }
 
    private void LoadZip()
    {
        if (_isLoaded)
            return;
 
        if (!File.Exists(FilePath))
        {
            _isLoaded = true;
            return;
        }
 
        using var stream = File.OpenRead(FilePath);
        LoadZip(stream);
    }
 
    private void LoadZip(Stream stream)
    {
        if (_isLoaded)
            return;
 
        try
        {
            Engine.Logger.LogInformation("Loading assets file.");
 
            _zipFile?.Close();
            _zipFile = null;
            _zipEntries.Clear();
 
            // SharpZipLib's ZipFile expects seekable input. Buffering here also gives
            // stream-loaded GAFs the same lifetime semantics as path-loaded GAFs:
            // the source stream is needed only for the duration of this call.
            MemoryStream? bufferedStream = stream.CanSeek ? null : new MemoryStream();
            Stream archiveStream = stream;
            if (bufferedStream is not null)
            {
                stream.CopyTo(bufferedStream);
                bufferedStream.Position = 0;
                archiveStream = bufferedStream;
            }
 
            _zipFile = new ZipFile(archiveStream) { IsStreamOwner = false };
 
            if (!string.IsNullOrEmpty(Password))
                _zipFile.Password = Password;
 
            // Force decryption on the first entry so a bad password fails early.
            var testEntry = _zipFile.Cast<ZipEntry>().FirstOrDefault(e => e.IsFile);
            if (testEntry != null)
            {
                using var testStream = _zipFile.GetInputStream(testEntry);
                testStream.ReadByte();
            }
 
            foreach (ZipEntry entry in _zipFile)
            {
                if (!entry.IsFile)
                    continue;
 
                var key = AssetsFileEntry.FromString(entry.Name);
                if (key == null)
                {
                    Engine.Logger.LogWarning("Invalid entry in assets file: {EntryName}", entry.Name);
                    continue;
                }
 
                using var entryStream = _zipFile.GetInputStream(entry);
                using var ms = new MemoryStream();
                entryStream.CopyTo(ms);
 
                _zipEntries[key] = ms.ToArray();
            }
 
            _isLoaded = true;
        }
        catch (ZipException zex)
        {
            Engine.Logger.LogError(zex, "Decryption failed — check password or file integrity.");
            throw;
        }
        catch (Exception ex)
        {
            Engine.Logger.LogError(ex, "Failed to load asset file.");
            throw;
        }
        finally
        {
            // Once contents are buffered, no need to keep the archive or source stream open.
            _zipFile?.Close();
            _zipFile = null;
        }
    }

    public void Add(AssetTypes type, string name, Func<Stream> streamFactory)
    {
        if (streamFactory == null)
            throw new ArgumentNullException(nameof(streamFactory));
 
        EnsureLoaded();
 
        using var stream = streamFactory();
        Add(type, name, stream);
    }

    public void Add(AssetTypes type, string filePath, string? name = null)
    {
        if (string.IsNullOrWhiteSpace(filePath))
            throw new ArgumentException("File path cannot be null or empty.", nameof(filePath));
 
        var nameToUse = name ?? Path.GetFileName(filePath);
        using var stream = File.OpenRead(filePath);
        Add(type, nameToUse, stream);
    }

    public void Add(AssetTypes type, string name, Stream stream)
    {
        if (stream == null)
            throw new ArgumentNullException(nameof(stream));
 
        if (string.IsNullOrWhiteSpace(name))
            throw new ArgumentException("Asset name cannot be null or empty.", nameof(name));
 
        if (!stream.CanRead)
            throw new ArgumentException("Stream must be readable.", nameof(stream));
 
        EnsureLoaded();
 
        var key = new AssetsFileEntry
        {
            AssetType = type,
            AssetName = name
        };
 
        using var ms = new MemoryStream();
        stream.CopyTo(ms);
        _zipEntries[key] = ms.ToArray();
    }

    public void Remove(AssetTypes type, string name)
    {
        if (string.IsNullOrWhiteSpace(name))
            throw new ArgumentException("Asset name cannot be null or empty.", nameof(name));
 
        EnsureLoaded();
 
        var key = new AssetsFileEntry
        {
            AssetType = type,
            AssetName = name
        };
 
        _zipEntries.Remove(key);
    }

    public Stream? this[AssetTypes type, string name] => Get(type, name);

    public Stream? Get(AssetTypes type, string name)
    {
        if (string.IsNullOrWhiteSpace(name))
            return null;
 
        EnsureLoaded();
 
        var exactKey = new AssetsFileEntry
        {
            AssetType = type,
            AssetName = name
        };
 
        if (_zipEntries.TryGetValue(exactKey, out var exactData))
            return new MemoryStream(exactData, writable: false);
 
        var requestedBaseName = Path.GetFileNameWithoutExtension(name);
 
        var match = _zipEntries.Keys.FirstOrDefault(k =>
            k.AssetType == type &&
            string.Equals(
                Path.GetFileNameWithoutExtension(k.AssetName),
                requestedBaseName,
                StringComparison.OrdinalIgnoreCase));
 
        return match != null && _zipEntries.TryGetValue(match, out var fallbackData)
            ? new MemoryStream(fallbackData, writable: false)
            : null;
    }

    public IEnumerable<AssetsFileEntry> GetAllEntries()
    {
        EnsureLoaded();
        return _zipEntries.Keys.ToList();
    }

    public void Save()
    {
        EnsureLoaded();
 
        _zipFile?.Close();
        _zipFile = null;
        _isLoaded = false;
 
        using (var fs = File.Create(FilePath))
        using (var zipStream = new ZipOutputStream(fs) { IsStreamOwner = true })
        {
            if (!string.IsNullOrEmpty(Password))
                zipStream.Password = Password;
 
            foreach (var (key, data) in _zipEntries)
            {
                var entry = new ZipEntry(key.ToString())
                {
                    DateTime = DateTime.Now
                };
 
                if (UseEncryption && !string.IsNullOrEmpty(Password))
                    entry.AESKeySize = 256;
 
                zipStream.PutNextEntry(entry);
 
                using var inputStream = new MemoryStream(data, writable: false);
                inputStream.CopyTo(zipStream);
                zipStream.CloseEntry();
            }
 
            zipStream.Finish();
        }
 
        Engine.Logger.LogInformation(
            "Assets file saved (Encrypted: {Encrypted}).",
            UseEncryption);
 
        // Keep the in-memory copy as the source of truth.
        _isLoaded = true;
    }

    public void Dispose()
    {
        _zipFile?.Close();
        _zipFile = null;
        _isLoaded = false;
        _zipEntries.Clear();
        _allAssetsFiles.Remove(this);
    }
}