Res2tf/Bson.cs

using System.Runtime.InteropServices;
using System.Text;
using System.Text.Json.Nodes;

namespace Res2tf;

using List = List<BsonValue>;
using Dict = Dictionary<string, BsonValue>;

public enum ElementType : byte
{
    Double = 1,
    String,
    Object,
    Array,
    Binary,
    Undefined,   // deprecated
    ObjectId,    // unsupported
    Boolean,
    DateTime,
    Null,
    RegExp,      // unsupported
    DbPointer,   // deprecated
    JsCode,      // unsupported
    Symbol,      // deprecated
    JsCodeScope, // deprecated
    Int32,
    Timestamp,   // unsupported
    Int64,
    Decimal128,  // unsupported
};

public class BsonValue : IEquatable<BsonValue>
{
    public static readonly BsonValue Null = new(ElementType.Null, null);

    public ElementType ElementType { get; }
    public object? Value { get; }

    public int Count => AsCollection().Count;
    public BsonValue this[int index]  { get => AsList()[index]; set => AsList()[index] = value; }
    public BsonValue this[string key] { get => AsDict()[key  ]; set => AsDict()[key  ] = value; }
    public Dict.KeyCollection   Keys   => AsDict().Keys;
    public Dict.ValueCollection Values => AsDict().Values;

    private BsonValue(ElementType type, object? value) { ElementType = type; Value = value; }

    public BsonValue(double   value) : this(ElementType.Double  , value) { }
    public BsonValue(string   value) : this(ElementType.String  , value) { ArgumentNullException.ThrowIfNull(value); }
    public BsonValue(byte[]   value) : this(ElementType.Binary  , value) { ArgumentNullException.ThrowIfNull(value); }
    public BsonValue(bool     value) : this(ElementType.Boolean , value) { }
    public BsonValue(DateTime value) : this(ElementType.DateTime, value) { }
    public BsonValue(int      value) : this(ElementType.Int32   , value) { }
    public BsonValue(long     value) : this(ElementType.Int64   , value) { }

    public static BsonValue Array() => new(ElementType.Array, new List());
    public static BsonValue Array(params IEnumerable<BsonValue> values) => new(ElementType.Array, values.ToList());

    public static BsonValue Object() => new(ElementType.Object, new Dict());
    public static BsonValue Object(IDictionary<string, BsonValue> entries) => new(ElementType.Object, entries.ToDictionary());
    public static BsonValue Object(params IEnumerable<(string, BsonValue)> entries) => new(ElementType.Object, entries.ToDictionary());

    public List AsList() => (Value as List) ?? throw new InvalidOperationException($"Attempting to access {ElementType} as Array");
    public Dict AsDict() => (Value as Dict) ?? throw new InvalidOperationException($"Attempting to access {ElementType} as Object");
    protected ICollection<BsonValue> AsCollection() => Value switch { List list => list, Dict dict => dict.Values, _ => throw new NotSupportedException($"Attempting to access {ElementType} as Array or Object"), };

    public JsonNode? ToJson()
        => Value switch {
            Dict     v => new JsonObject(AsDict().Select(e => new KeyValuePair<string, JsonNode?>(e.Key, e.Value.ToJson()))),
            List     v => new JsonArray (AsList().Select(e => e.ToJson()).ToArray()),
            double   v => JsonValue.Create(v),
            string   v => JsonValue.Create(v)!,
            byte[]   v => JsonValue.Create(v)!,
            bool     v => JsonValue.Create(v),
            DateTime v => JsonValue.Create(v),
            int      v => JsonValue.Create(v),
            long     v => JsonValue.Create(v),
            null       => null,
                     _ => throw new NotImplementedException(),
        };

    public T[] ToArray<T>()
    {
        if (typeof(T) == typeof(BsonValue)) return (T[])(object)AsList().ToArray();
        if (typeof(T) == typeof(byte) && ElementType == ElementType.Binary) return (T[])Value!;
        if (typeof(T) == typeof(float)) return AsList().Select(v => (T)(object)(float)(double)v.Value!).ToArray();
        if (ElementTypeOf<T>() == null) throw new NotSupportedException($"{typeof(T)} is not a valid ElementType");
        return AsList().Select(v => (T)v.Value!).ToArray();
    }

    public Dictionary<string, T> ToDictionary<T>()
    {
        if (typeof(T) == typeof(BsonValue)) return (Dictionary<string, T>)(object)AsDict().ToDictionary();
        if (ElementTypeOf<T>() == null) throw new NotSupportedException($"{typeof(T)} is not a valid ElementType");
        return AsDict().ToDictionary(e => e.Key, e => (T)e.Value.Value!);
    }

    public BsonValue? GetOrNull(string key) => AsDict().GetValueOrDefault(key);

    public void Add(BsonValue value)             => AsList().Add(value);
    public void Add(string key, BsonValue value) => AsDict().Add(key, value);

    public bool Contains(BsonValue value) => AsList().Contains(value);
    public bool ContainsKey(string key)   => AsDict().ContainsKey(key);

    public void RemoveAt(int index) => AsList().RemoveAt(index);
    public bool Remove(string key)  => AsDict().Remove(key);

    public void Clear()
    {
        switch (Value) {
            case List list: list.Clear(); break;
            case Dict dict: dict.Clear(); break;
            default: throw new NotSupportedException($"Attempting to access {ElementType} as Array or Object");
        };
    }

    public override bool Equals(object? obj)
        => Equals(obj as BsonValue);
    public bool Equals(BsonValue? other)
        => (other is not null)
        && (ElementType == other.ElementType)
        && ElementType switch {
            ElementType.Array  => Enumerable.SequenceEqual(AsList(), other.AsList()),
            ElementType.Object => DictionaryEquals(AsDict(), other.AsDict()),
            _ => Equals(Value, other.Value),
        };
    public override int GetHashCode()
        => HashCode.Combine(ElementType, Value);

    public static bool operator ==(BsonValue? left, BsonValue? right)
        => ReferenceEquals(left, right) || (left?.Equals(right) ?? false);
    public static bool operator !=(BsonValue? left, BsonValue? right)
        => !(left == right);

    public static implicit operator BsonValue(double   value) => new(value);
    public static implicit operator BsonValue(string   value) => new(value);
    public static implicit operator BsonValue(byte[]   value) => new(value);
    public static implicit operator BsonValue(bool     value) => new(value);
    public static implicit operator BsonValue(DateTime value) => new(value);
    public static implicit operator BsonValue(int      value) => new(value);
    public static implicit operator BsonValue(long     value) => new(value);

    public static explicit operator double  (BsonValue value) => (double  )value.Value!;
    public static explicit operator string  (BsonValue value) => (string  )value.Value!;
    public static explicit operator byte[]  (BsonValue value) => (byte[]  )value.Value!;
    public static explicit operator bool    (BsonValue value) => (bool    )value.Value!;
    public static explicit operator DateTime(BsonValue value) => (DateTime)value.Value!;
    public static explicit operator int     (BsonValue value) => (int     )value.Value!;
    public static explicit operator long    (BsonValue value) => (long    )value.Value!;

    public static explicit operator float(BsonValue value) => (float)(double)value;

    private static ElementType? ElementTypeOf<T>()
    {
        if (typeof(T) == typeof(double  )) return ElementType.Double;
        if (typeof(T) == typeof(string  )) return ElementType.String;
        if (typeof(T) == typeof(Dict    )) return ElementType.Object;
        if (typeof(T) == typeof(List    )) return ElementType.Array;
        if (typeof(T) == typeof(byte[]  )) return ElementType.Binary;
        if (typeof(T) == typeof(bool    )) return ElementType.Boolean;
        if (typeof(T) == typeof(DateTime)) return ElementType.DateTime;
        if (typeof(T) == typeof(int     )) return ElementType.Int32;
        if (typeof(T) == typeof(long    )) return ElementType.Int64;
        return null;
    }

    private static bool DictionaryEquals(Dict first, Dict second)
        => (first.Count == second.Count) && !first.Except(second).Any();
}

public class BsonReader : IDisposable
{
    private readonly BinaryReader _reader;
    private readonly List<byte> _cStringBuf = [];

    public BsonReader(Stream stream, bool leaveOpen = false)
    {
        if (!stream.CanRead) throw new ArgumentException("stream is not readable");
        _reader = new BinaryReader(stream, Encoding.UTF8, leaveOpen);
    }

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

    public static BsonValue Load(Stream stream, bool leaveOpen = false)
    {
        using var reader = new BsonReader(stream, leaveOpen);
        return reader.ReadObject();
    }
    public static BsonValue Load(byte[] bytes)
    {
        using var stream = new MemoryStream(bytes);
        return Load(stream);
    }

    public BsonValue ReadObject()
    {
        // A "document" is prefixed with the number of bytes it takes up, but
        // we can just ignore it and read elements until we detect a NUL byte.
        // Thankfully we don't need to write a BsonWriter ...
        _reader.ReadInt32(); // length, unused

        var result = BsonValue.Object();
        while (ReadElement() is (string name, BsonValue value))
            result.Add(name, value);
        return result;
    }

    public BsonValue ReadArray()
    {
        // An "array" is just an object with stringified integer keys ...
        _reader.ReadInt32(); // length, unused

        var result = BsonValue.Array();
        while (ReadElement() is (string name, BsonValue value))
        {
            if (!int.TryParse(name, out var index)) throw new Exception("Couldn't parse array index as integer");
            if (index != result.Count) throw new Exception("Non-sequential array index");
            result.Add(value);
        }
        return result;
    }

    public BsonValue ReadBinary()
    {
        var length = _reader.ReadInt32();
        _reader.ReadByte(); // subtype, unused
        return _reader.ReadBytes(length);
    }

    public (string Name, BsonValue Value)? ReadElement()
    {
        var type = (ElementType)_reader.ReadByte();
        if (type == 0) return null;

        var name = ReadCString();
        return (name, type switch
        {
            ElementType.Double   => _reader.ReadDouble(),
            ElementType.String   => ReadLengthPrefixedString(),
            ElementType.Object   => ReadObject(),
            ElementType.Array    => ReadArray(),
            ElementType.Binary   => ReadBinary(),
            ElementType.Boolean  => _reader.ReadByte() != 0,
            ElementType.DateTime => DateTime.UnixEpoch + TimeSpan.FromMilliseconds(_reader.ReadInt64()),
            ElementType.Null     => BsonValue.Null,
            ElementType.Int32    => _reader.ReadInt32(),
            ElementType.Int64    => _reader.ReadInt64(),

            ElementType.Undefined | ElementType.ObjectId | ElementType.RegExp
                | ElementType.DbPointer | ElementType.JsCode | ElementType.Symbol
                | ElementType.JsCodeScope | ElementType.Timestamp | ElementType.Decimal128
                => throw new NotSupportedException($"Reading {type} is not supported"),

            _ => throw new Exception($"Unknown ElementType {type}"),
        });
    }

    private string ReadLengthPrefixedString()
    {
        var length = _reader.ReadInt32();
        var bytes = _reader.ReadBytes(length);
        if (bytes[^1] != 0) throw new Exception("Trailing byte in string must be NUL");
        return Encoding.UTF8.GetString(bytes.AsSpan()[0..^1]);
    }

    private string ReadCString()
    {
        _cStringBuf.Clear();
        byte b; while ((b = _reader.ReadByte()) != 0) _cStringBuf.Add(b);
        var bytes = CollectionsMarshal.AsSpan(_cStringBuf);
        return Encoding.UTF8.GetString(bytes);
    }
}