ReFuel.Gltf/ReFuel.Gltf/GltfDocument.cs

using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.Text.Json;
using System.Text.Json.Serialization;
using ReFuel.Gltf.IO;

namespace ReFuel.Gltf
{
    [Flags]
    public enum GltfDocumentKind
    {
        Json = 1 << 0,
        Binary = 1 << 1,
        Any = Json | Binary,
    }


    /// <summary>
    /// glTF Document Object GLObjects
    /// </summary>
    [JsonConverter(typeof(GltfJsonConverter))]
    public class GltfDocument : GltfObject, IDisposable
    {
        public override GltfObjectKind Kind => GltfObjectKind.Document;

        public List<Stream> BinaryStreams { get; private set; }

        public int DefaultScene { get; set; } = 0;
        public HashSet<string> Extensions { get; set; } = new HashSet<string>();
        public HashSet<string> Required { get; set; } = new HashSet<string>();
        public GltfAsset Asset { get; }
        public GltfIdList<GltfScene> Scenes { get; } = new GltfIdList<GltfScene>();
        public GltfIdList<GltfNode> Nodes { get; } = new GltfIdList<GltfNode>();
        public GltfIdList<GltfMesh> Meshes { get; } = new GltfIdList<GltfMesh>();
        public GltfIdList<GltfBuffer> Buffers { get; } = new GltfIdList<GltfBuffer>();
        public GltfIdList<GltfBufferView> BufferViews { get; } = new GltfIdList<GltfBufferView>();
        public GltfIdList<GltfAccessor> Accessors { get; } = new GltfIdList<GltfAccessor>();
        public GltfIdList<GltfMaterial> Materials { get; } = new GltfIdList<GltfMaterial>();
        public GltfIdList<GltfTexture> Textures { get; } = new GltfIdList<GltfTexture>();
        public GltfIdList<GltfSampler> Samplers { get; } = new GltfIdList<GltfSampler>();
        public GltfIdList<GltfImage> Images { get; } = new GltfIdList<GltfImage>();
        public GltfIdList<GltfSkin> Skins { get; } = new GltfIdList<GltfSkin>();

        private GltfDocument(JsonDocument doc, IEnumerable<StreamWrapper>? binaryStreams = null)
        {
            Document = this;
            Asset = new GltfAsset(this);

            if (binaryStreams == null)
                BinaryStreams = new List<Stream>();
            else
                BinaryStreams = new List<Stream>(binaryStreams);

            JsonElement root = doc.RootElement;
            Deserialize(root);
        }

        private static void WriteStringSet(HashSet<string> set, JsonElement stringArray)
        {
            if (stringArray.ValueKind != JsonValueKind.Array)
                throw new Exception("Expected a string array.");

            foreach (JsonElement item in stringArray.EnumerateArray())
            {
                string? str = item.GetString();
                if (str != null)
                    set.Add(str);
            }
        }

        public void Save(Stream str, GltfDocumentKind kind)
        {
            if (BinaryStreams.Count > 0)
            {
                if (kind == GltfDocumentKind.Json)
                    throw new NotSupportedException("Cannot write binary streams into a JSON style glTF document.");
                else if (kind == GltfDocumentKind.Any)
                    kind = GltfDocumentKind.Binary;
            }

            switch (kind)
            {
                case GltfDocumentKind.Json:
                    SaveJson(str);
                    break;
                default:
                case GltfDocumentKind.Binary:
                    SaveBinary(str);
                    break;
            }
        }

        private void SaveJson(Stream str)
        {
            throw new NotImplementedException();
        }

        private void SaveBinary(Stream str)
        {
            throw new NotImplementedException();
        }

        private static GltfDocument OpenJson(StreamWrapper str)
        {
            JsonDocument doc = JsonDocument.Parse(str, new JsonDocumentOptions()
            {
                AllowTrailingCommas = true,
                CommentHandling = JsonCommentHandling.Skip
            });

            return new GltfDocument(doc);
        }

        internal static GltfDocument OpenJson(ref Utf8JsonReader reader)
        {
            if (!JsonDocument.TryParseValue(ref reader, out JsonDocument? document))
            {
                throw new JsonException();
            }

            return new GltfDocument(document);
        }

        private static GltfDocument OpenBinary(StreamWrapper str)
        {
            if (GlbHeader.Read(str, out GlbHeader header))
            {
                return OpenBinary(str, header);
            }

            throw new Exception("File does not have a valid GLB header.");
        }

        private static GltfDocument OpenBinary(StreamWrapper str, in GlbHeader header)
        {
            if (header.Version < 2)
            {
                throw new NotSupportedException("Only glTF version 2.0 files are supported.");
            }

            JsonDocument? document = null;
            List<StreamWrapper> binaryStreams = new List<StreamWrapper>();

            while (GlbChunk.Read(str, out GlbChunk chunk))
            {
                using StreamWrapper substream = str.ForkAbs(str.Position, str.Position + chunk.Size, true);

                if (chunk.Type == GlbChunkType.Json)
                {
                    document = JsonDocument.Parse(substream.Fork(keepOpen:true));
                }
                else
                {
                    Stream copy = new MemoryStream();

                    substream.CopyTo(copy);
                    copy.Position = 0;

                    binaryStreams.Add(new StreamWrapper(copy));
                }
            }

            if (document == null)
            {
                throw new Exception("This glTF document has no manifest JSON.");
            }
            else
            {
                return new GltfDocument(document, binaryStreams);
            }
        }

        public static GltfDocument Open(Stream str, GltfDocumentKind kind = GltfDocumentKind.Any)
        {
            using StreamWrapper wrapper = new StreamWrapper(str, true);

            switch (kind)
            {
                case GltfDocumentKind.Json:
                    return OpenJson(wrapper);
                case GltfDocumentKind.Binary:
                    return OpenBinary(wrapper);
                default:
                case GltfDocumentKind.Any:
                    if (GlbHeader.Read(wrapper, out GlbHeader header))
                    {
                        return OpenBinary(wrapper, header);
                    }
                    else
                    {
                        str.Seek(0, SeekOrigin.Begin);
                        return OpenJson(wrapper);
                    }
            }
        }

        private bool isDisposed = false;

        public void Dispose() => Dispose(true);
        private void Dispose(bool disposing)
        {
            if (isDisposed) return;

            if (disposing)
            {
                foreach (Stream str in BinaryStreams)
                    str.Dispose();

                GC.SuppressFinalize(this);
            }

            isDisposed = true;
        }

        private static Dictionary<string, IGltfExtensionProvider> extensionProviders = new Dictionary<string, IGltfExtensionProvider>();
        private static HashSet<string> supportedExtensions = new HashSet<string>();

        public static IReadOnlySet<string> SupportedExtensions => supportedExtensions;

        static GltfDocument()
        {
            // AddExtension(RfChecksumProvider.ExtensionProviders);
            // AddExtension(RfCompression.ExtensionProviders);
        }

        public static void AddExtension(IGltfExtensionProvider provider)
        {
            supportedExtensions.Add(provider.Name);
            extensionProviders.Add(provider.Name, provider);
        }

        public static void AddExtension<T>() where T : IGltfExtensionProvider, new()
        {
            AddExtension(new T());
        }

        public static void AddExtension(IEnumerable<IGltfExtensionProvider> providers)
        {
            foreach (var provider in providers) AddExtension(provider);
        }

        public static bool TryGetExtensionProvider(string name, [NotNullWhen(true)] out IGltfExtensionProvider? provider)
        {
            return extensionProviders.TryGetValue(name, out provider);
        }

        internal override void Deserialize(JsonElement root)
        {
            if (root.ValueKind != JsonValueKind.Object)
                throw new Exception("The glTF JSON document has an object as its root element");

            Dictionary<string, JsonElement> elements = new Dictionary<string, JsonElement>();
            foreach (JsonProperty property in root.EnumerateObject())
            {
                // Some elements may be interpreted earlier than others.
                // This is mainly the default child, asset information, and the extensions.
                switch (property.Name)
                {
                    case "asset":
                        Asset.Deserialize(property.Value);
                        if (Asset.MinVersion > new Version(2, 0) || Asset.Version != new Version(2, 0))
                            throw new Exception($"Unsupported glTF version {Asset.MinVersion ?? Asset.Version}");
                        break;

                    case "extensionsUsed":
                        WriteStringSet(Extensions, property.Value);
                        break;

                    case "requiredExtensions":
                        WriteStringSet(Required, property.Value);
                        break;

                    case "child":
                        if (property.Value.TryGetInt32(out int defaultScene))
                            DefaultScene = defaultScene;
                        break;
                    default:
                        elements.Add(property.Name, property.Value);
                        break;
                }
            }

            foreach (string value in Required)
            {
                if (!SupportedExtensions.Contains(value))
                {
                    throw new Exception($"Need unsupported extension {value} for this glTF file.");
                }
            }

            foreach ((string name, JsonElement element) in elements)
            {
                switch (name)
                {
                    case "extension": Extension.Parse(element); break;
                    case "extra": Extras.Parse(element); break;
                    case "scenes":
                        AssertArray(element, "scenes");
                        foreach (JsonElement child in element.EnumerateArray())
                        {
                            GltfScene scene = new GltfScene(this);
                            scene.Deserialize(child);
                            Scenes.Append(scene);
                        }
                        break;
                    case "nodes":
                        AssertArray(element, "nodes");
                        foreach (JsonElement child in element.EnumerateArray())
                        {
                            GltfNode node = new GltfNode(this);
                            node.Deserialize(child);
                            Nodes.Append(node);
                        }
                        break;
                    case "buffers":
                        AssertArray(element, "buffers");
                        foreach (JsonElement child in element.EnumerateArray())
                        {
                            GltfBuffer buffer = new GltfBuffer(this);
                            buffer.Deserialize(child);
                            Buffers.Append(buffer);
                        }
                        break;
                    case "bufferViews":
                        AssertArray(element, "bufferViews");
                        foreach (JsonElement child in element.EnumerateArray())
                        {
                            GltfBufferView view = new GltfBufferView(this);
                            view.Deserialize(child);
                            BufferViews.Append(view);
                        }
                        break;
                    case "accessors":
                        AssertArray(element, "accessors");
                        foreach (JsonElement child in element.EnumerateArray())
                        {
                            GltfAccessor accessor = new GltfAccessor(this);
                            accessor.Deserialize(child);
                            Accessors.Append(accessor);
                        }
                        break;
                    case "meshes":
                        AssertArray(element, "meshes");
                        foreach (JsonElement child in element.EnumerateArray())
                        {
                            GltfMesh mesh = new GltfMesh(this);
                            mesh.Deserialize(child);
                            Meshes.Append(mesh);
                        }
                        break;
                    case "materials":
                        AssertArray(element, "materials");
                        foreach (JsonElement child in element.EnumerateArray())
                        {
                            GltfMaterial material = new GltfMaterial(this);
                            material.Deserialize(child);
                            Materials.Append(material);
                        }
                        break;
                    case "textures":
                        AssertArray(element, "textures");
                        foreach (JsonElement child in element.EnumerateArray())
                        {
                            GltfTexture texture = new GltfTexture(this);
                            texture.Deserialize(child);
                            Textures.Append(texture);
                        }
                        break;
                    case "samplers":
                        AssertArray(element, "samplers");
                        foreach (JsonElement child in element.EnumerateArray())
                        {
                            GltfSampler sampler = new GltfSampler(this);
                            sampler.Deserialize(child);
                            Samplers.Append(sampler);
                        }
                        break;
                    case "images":
                        AssertArray(element, "images");
                        foreach (JsonElement child in element.EnumerateArray())
                        {
                            GltfImage image = new GltfImage(this);
                            image.Deserialize(child);
                            Images.Append(image);
                        }
                        break;
                    case "skins":
                        AssertArray(element, "skins");
                        foreach (JsonElement child in element.EnumerateArray())
                        {
                            GltfSkin skin = new GltfSkin(this);
                            skin.Deserialize(element);
                            Skins.Append(skin);
                        }
                        break;
                }
            }
        }

        internal override void Serialize(Utf8JsonWriter writer)
        {
            throw new NotImplementedException();
        }
    }
}