Res2tf/ResonitePackage.cs

using System.Diagnostics.CodeAnalysis;
using System.IO.Compression;
using System.Numerics;
using System.Text.Json;
using Elements.Assets;
using SharpGLTF.Materials;
using SharpGLTF.Schema2;

namespace Res2tf;

public class ResonitePackage : IDisposable
{
    readonly ZipArchive _zip;

    public DataTree Tree { get; }
    public BsonValue Bson => Tree.Bson;
    public Dictionary<string, object> Assets { get; } = [];

    public ResonitePackage(string path)
    {
        _zip = new ZipArchive(File.OpenRead(path));

        using var mainRecord = _zip.GetEntry("R-Main.record")!.Open();
        var mainDocument = JsonDocument.Parse(mainRecord);
        var mainAssetUri = mainDocument.RootElement.GetProperty("assetUri").GetString()!;

        Tree = new(ResolveToStream(mainAssetUri));
    }

    public void Dispose()
    {
        _zip.Dispose();
        GC.SuppressFinalize(this);
    }

    public byte[] ResolveToBytes(string uriString)
    {
        using var stream = ResolveToStream(uriString, true);
        return ((MemoryStream)stream).ToArray();
    }
    public Stream ResolveToStream(string uriString, bool copy = false)
    {
        // URIs coming from the DataTree are gonna start with an '@' symbol, so just strip it.
        if (uriString.StartsWith('@')) uriString = uriString[1..];

        var uri = new Uri(uriString);
        if (uri.Scheme != "packdb") throw new ArgumentException("Not a packdb URI");
        if (!uri.AbsolutePath.StartsWith('/')) throw new ArgumentException("Not valid packdb URI");

        var stream = _zip.GetEntry($"Assets{uri.AbsolutePath}")!.Open();
        if (!copy) return stream;
        else using (stream) {
            var bytes = new MemoryStream();
            stream.CopyTo(bytes);
            bytes.Seek(0, SeekOrigin.Begin);
            return bytes;
        }
    }

    public void ProcessAssets(ModelRoot model)
    {
        // Process static assets first.
        var staticMesh    = Tree.GetTypeId("[FrooxEngine]FrooxEngine.StaticMesh");
        var staticTexture = Tree.GetTypeId("[FrooxEngine]FrooxEngine.StaticTexture2D");
        foreach (var asset in Bson["Assets"].AsList()) {
            var type = (int)asset["Type"];

            if (type == staticMesh) {
                var id   = (string)asset["Data"]["ID"];
                var uri  = (string)asset["Data"]["URL"]["Data"];
                var mesh = new MeshAsset(new MeshX(ResolveToStream(uri, true)), model);
                Assets[id] = mesh;
            }

            else if (type == staticTexture) {
                var id  = (string)asset["Data"]["ID"];
                var uri = (string)asset["Data"]["URL"]["Data"];
                // For now, texture "asset" is just its URI, material will load it from there.
                Assets[id] = uri;
            }
        }

        var toonMaterial  = Tree.GetTypeId("[FrooxEngine]FrooxEngine.XiexeToonMaterial");
        var unlitMaterial = Tree.GetTypeId("[FrooxEngine]FrooxEngine.UnlitMaterial");
        foreach (var asset in Bson["Assets"].AsList()) {
            var type = (int)asset["Type"];

            if (type == toonMaterial) {
                var id = (string)asset["Data"]["ID"];

                var color = ResoniteToColor(asset["Data"]["Color"]);

                var textureId   = (string?)asset["Data"]["MainTexture"]["Data"];
                var texturePath = (textureId   != null) ? (string)Assets[textureId]   : null;
                var texture     = (texturePath != null) ? ResolveToBytes(texturePath) : null;

                var material = new MaterialBuilder()
                    .WithDoubleSide(false)
                    .WithBaseColor(color)
                    .WithChannelImage(KnownChannel.BaseColor, texture);
                Assets[id] = material;
            }

            else if (type == unlitMaterial) {
                var id = (string)asset["Data"]["ID"];
    
                var color = ResoniteToColor(asset["Data"]["TintColor"]);

                var textureId   = (string)asset["Data"]["Texture"]["Data"];
                var texturePath = (textureId   != null) ? (string)Assets[textureId]   : null;
                var texture     = (texturePath != null) ? ResolveToBytes(texturePath) : null;

                var material = new MaterialBuilder()
                    .WithDoubleSide(false)
                    .WithBaseColor(color)
                    .WithChannelImage(KnownChannel.BaseColor, texture);
                Assets[id] = material;
            }
        }
    }

    public bool TryGetAsset<TAsset>(string assetId, [NotNullWhen(true)] out TAsset? asset)
    {
        if (Assets.TryGetValue(assetId, out var assetUntyped) && (assetUntyped is TAsset a))
            { asset = a; return true; } else { asset = default; return false; }
    }

    static Vector4 ResoniteToColor(BsonValue bson)
    {
        var list = bson["Data"].AsList();
        var r = Math.Clamp((float)list[0], 0.0f, 1.0f);
        var g = Math.Clamp((float)list[1], 0.0f, 1.0f);
        var b = Math.Clamp((float)list[2], 0.0f, 1.0f);
        var a = Math.Clamp((float)list[3], 0.0f, 1.0f);
        return new(r, g, b, a);
    }
}