using System.IO.Compression;
using System.Numerics;
using System.Text;

namespace Res2tf;

public class DataTree
{
    public BsonValue Bson { get; }
    public string[] Types { get; }
    public DataTreeNode Root { get; }

    // Technically unused, but kept in case it will be needed.
    public Dictionary<string, DataTreeNode> Nodes { get; } = [];

    public DataTree(string path)   : this(LoadBson(path)) {  }
    public DataTree(Stream stream) : this(LoadBson(stream)) {  }
    public DataTree(BsonValue bson)
    {
        Bson  = bson;
        Types = bson["Types"].ToArray<string>();
        Root  = CreateNodeRecursive(null, bson["Object"]);
    }

    public int? GetTypeId(string type)
        => Array.IndexOf(Types, type) switch
            { -1 => null, int i => i };

    DataTreeNode CreateNodeRecursive(DataTreeNode? parent, BsonValue bson)
    {
        var node = new DataTreeNode(this, parent, bson);
        foreach (var childBson in bson.GetOrNull("Children")?.AsList() ?? [])
            node.Children.Add(CreateNodeRecursive(node, childBson));
        Nodes[node.ID] = node;
        return node;
    }

    static BsonValue LoadBson(string path)
        => LoadBson(File.OpenRead(path));
    static BsonValue LoadBson(Stream stream)
    {
        // Read and verify file header.
        using var reader = new BinaryReader(stream);
        var magic       = Encoding.ASCII.GetString(reader.ReadBytes(4));
        var version     = reader.ReadInt32();
        var compression = reader.ReadByte();
        if (magic  != "FrDT") throw new Exception("Invalid magic number");
        if (version     != 0) throw new Exception("Unknown version");
        if (compression != 3) throw new Exception("Unsupported compression");

        // Decompress and parse BSON.
        using var brotli = new BrotliStream(stream, CompressionMode.Decompress);
        return BsonReader.Load(brotli);
    }
}

public record DataTreeNode(DataTree Tree, DataTreeNode? Parent, BsonValue Bson)
{
    public List<DataTreeNode> Children { get; } = [];

    public string ID   => (string)Bson["ID"];
    public string Name => (string)Bson["Name"]["Data"];

    public Vector3    Position { get { var p = Bson["Position"]["Data"].ToArray<float>(); return new(p[0], p[1], p[2]); } }
    public Quaternion Rotation { get { var r = Bson["Rotation"]["Data"].ToArray<float>(); return new(r[0], r[1], r[2], r[3]); } }
    public Vector3    Scale    { get { var s = Bson["Scale"]   ["Data"].ToArray<float>(); return new(s[0], s[1], s[2]); } }

    public void Remove() {
        Tree.Nodes.Remove(ID);
        Parent?.Children.Remove(this);
    }

    public IEnumerable<DataTreeNode> EnumerateRecursive()
    {
        yield return this;
        // Iterate children in reverse order so calling Remove on the current child is safe?
        // for (var i = Children.Count - 1; i >= 0; i--)
        foreach (var child in Children)
            foreach (var elem in child.EnumerateRecursive())
                yield return elem;
    }
}