using System;

namespace ZB.MOM.WW.CBDD.Bson;

/// <summary>
/// Represents an in-memory BSON document with lazy parsing.
/// Uses Memory&lt;byte&gt; to store raw BSON data for zero-copy operations.
/// </summary>
public sealed class BsonDocument
{
    private readonly Memory<byte> _rawData;
    private readonly System.Collections.Concurrent.ConcurrentDictionary<ushort, string>? _keys;

    /// <summary>
    /// Initializes a new instance of the <see cref="BsonDocument"/> class from raw BSON memory.
    /// </summary>
    /// <param name="rawBsonData">The raw BSON data.</param>
    /// <param name="keys">The optional key dictionary.</param>
    public BsonDocument(Memory<byte> rawBsonData, System.Collections.Concurrent.ConcurrentDictionary<ushort, string>? keys = null)
    {
        _rawData = rawBsonData;
        _keys = keys;
    }

    /// <summary>
    /// Initializes a new instance of the <see cref="BsonDocument"/> class from raw BSON bytes.
    /// </summary>
    /// <param name="rawBsonData">The raw BSON data.</param>
    /// <param name="keys">The optional key dictionary.</param>
    public BsonDocument(byte[] rawBsonData, System.Collections.Concurrent.ConcurrentDictionary<ushort, string>? keys = null)
    {
        _rawData = rawBsonData;
        _keys = keys;
    }

    /// <summary>
    /// Gets the raw BSON bytes
    /// </summary>
    public ReadOnlySpan<byte> RawData => _rawData.Span;

    /// <summary>
    /// Gets the document size in bytes
    /// </summary>
    public int Size => BitConverter.ToInt32(_rawData.Span[..4]);

    /// <summary>
    /// Creates a reader for this document
    /// </summary>
    public BsonSpanReader GetReader() => new BsonSpanReader(_rawData.Span, _keys ?? new System.Collections.Concurrent.ConcurrentDictionary<ushort, string>());

    /// <summary>
    /// Tries to get a field value by name.
    /// Returns false if field not found.
    /// </summary>
    /// <param name="fieldName">The field name.</param>
    /// <param name="value">When this method returns, contains the field value if found; otherwise <see langword="null"/>.</param>
    /// <returns><see langword="true"/> if the field is found; otherwise, <see langword="false"/>.</returns>
    public bool TryGetString(string fieldName, out string? value)
    {
        value = null;
        var reader = GetReader();
        if (reader.Remaining < 5)
            return false;

        _ = reader.ReadDocumentSize();
        fieldName = fieldName.ToLowerInvariant();
        while (reader.Remaining > 1)
        {
            var type = reader.ReadBsonType();
            if (type == BsonType.EndOfDocument)
                break;

            var name = reader.ReadElementHeader();

            if (name == fieldName && type == BsonType.String)
            {
                value = reader.ReadString();
                return true;
            }

            reader.SkipValue(type);
        }

        return false;
    }

    /// <summary>
    /// Tries to get an Int32 field value by name.
    /// </summary>
    /// <param name="fieldName">The field name.</param>
    /// <param name="value">When this method returns, contains the field value if found; otherwise zero.</param>
    /// <returns><see langword="true"/> if the field is found; otherwise, <see langword="false"/>.</returns>
    public bool TryGetInt32(string fieldName, out int value)
    {
        value = 0;
        var reader = GetReader();
        if (reader.Remaining < 5)
            return false;

        _ = reader.ReadDocumentSize();
        fieldName = fieldName.ToLowerInvariant();
        while (reader.Remaining > 1)
        {
            var type = reader.ReadBsonType();
            if (type == BsonType.EndOfDocument)
                break;

            var name = reader.ReadElementHeader();

            if (name == fieldName && type == BsonType.Int32)
            {
                value = reader.ReadInt32();
                return true;
            }

            reader.SkipValue(type);
        }

        return false;
    }

    /// <summary>
    /// Tries to get an ObjectId field value by name.
    /// </summary>
    /// <param name="fieldName">The field name.</param>
    /// <param name="value">When this method returns, contains the field value if found; otherwise default.</param>
    /// <returns><see langword="true"/> if the field is found; otherwise, <see langword="false"/>.</returns>
    public bool TryGetObjectId(string fieldName, out ObjectId value)
    {
        value = default;
        var reader = GetReader();
        if (reader.Remaining < 5)
            return false;

        _ = reader.ReadDocumentSize();
        fieldName = fieldName.ToLowerInvariant();
        while (reader.Remaining > 1)
        {
            var type = reader.ReadBsonType();
            if (type == BsonType.EndOfDocument)
                break;

            var name = reader.ReadElementHeader();

            if (name == fieldName && type == BsonType.ObjectId)
            {
                value = reader.ReadObjectId();
                return true;
            }

            reader.SkipValue(type);
        }

        return false;
    }

    /// <summary>
    /// Creates a new BsonDocument from field values using a builder pattern
    /// </summary>
    /// <param name="keyMap">The key map used for field name encoding.</param>
    /// <param name="buildAction">The action that populates the builder.</param>
    /// <returns>The created BSON document.</returns>
    public static BsonDocument Create(System.Collections.Concurrent.ConcurrentDictionary<string, ushort> keyMap, Action<BsonDocumentBuilder> buildAction)
    {
        var builder = new BsonDocumentBuilder(keyMap);
        buildAction(builder);
        return builder.Build();
    }
}

/// <summary>
/// Builder for creating BSON documents
/// </summary>
public sealed class BsonDocumentBuilder
{
    private byte[] _buffer = new byte[1024]; // Start with 1KB
    private int _position;
    private readonly System.Collections.Concurrent.ConcurrentDictionary<string, ushort> _keyMap;

    /// <summary>
    /// Initializes a new instance of the <see cref="BsonDocumentBuilder"/> class.
    /// </summary>
    /// <param name="keyMap">The key map used for field name encoding.</param>
    public BsonDocumentBuilder(System.Collections.Concurrent.ConcurrentDictionary<string, ushort> keyMap)
    {
        _keyMap = keyMap;
        var writer = new BsonSpanWriter(_buffer, _keyMap);
        _position = writer.Position;
    }

    /// <summary>
    /// Adds a string field to the document.
    /// </summary>
    /// <param name="name">The field name.</param>
    /// <param name="value">The field value.</param>
    /// <returns>The current builder.</returns>
    public BsonDocumentBuilder AddString(string name, string value)
    {
        EnsureCapacity(256); // Conservative estimate
        var writer = new BsonSpanWriter(_buffer.AsSpan(_position..), _keyMap);
        writer.WriteString(name, value);
        _position += writer.Position;
        return this;
    }

    /// <summary>
    /// Adds an Int32 field to the document.
    /// </summary>
    /// <param name="name">The field name.</param>
    /// <param name="value">The field value.</param>
    /// <returns>The current builder.</returns>
    public BsonDocumentBuilder AddInt32(string name, int value)
    {
        EnsureCapacity(64);
        var writer = new BsonSpanWriter(_buffer.AsSpan(_position..), _keyMap);
        writer.WriteInt32(name, value);
        _position += writer.Position;
        return this;
    }

    /// <summary>
    /// Adds an Int64 field to the document.
    /// </summary>
    /// <param name="name">The field name.</param>
    /// <param name="value">The field value.</param>
    /// <returns>The current builder.</returns>
    public BsonDocumentBuilder AddInt64(string name, long value)
    {
        EnsureCapacity(64);
        var writer = new BsonSpanWriter(_buffer.AsSpan(_position..), _keyMap);
        writer.WriteInt64(name, value);
        _position += writer.Position;
        return this;
    }

    /// <summary>
    /// Adds a Boolean field to the document.
    /// </summary>
    /// <param name="name">The field name.</param>
    /// <param name="value">The field value.</param>
    /// <returns>The current builder.</returns>
    public BsonDocumentBuilder AddBoolean(string name, bool value)
    {
        EnsureCapacity(64);
        var writer = new BsonSpanWriter(_buffer.AsSpan(_position..), _keyMap);
        writer.WriteBoolean(name, value);
        _position += writer.Position;
        return this;
    }

    /// <summary>
    /// Adds an ObjectId field to the document.
    /// </summary>
    /// <param name="name">The field name.</param>
    /// <param name="value">The field value.</param>
    /// <returns>The current builder.</returns>
    public BsonDocumentBuilder AddObjectId(string name, ObjectId value)
    {
        EnsureCapacity(64);
        var writer = new BsonSpanWriter(_buffer.AsSpan(_position..), _keyMap);
        writer.WriteObjectId(name, value);
        _position += writer.Position;
        return this;
    }

    /// <summary>
    /// Builds a BSON document from the accumulated fields.
    /// </summary>
    /// <returns>The constructed BSON document.</returns>
    public BsonDocument Build()
    {
        // Layout: [int32 size][field bytes...][0x00 terminator]
        var totalSize = _position + 5;
        var finalBuffer = new byte[totalSize];

        BitConverter.TryWriteBytes(finalBuffer.AsSpan(0, 4), totalSize);
        _buffer.AsSpan(0, _position).CopyTo(finalBuffer.AsSpan(4, _position));
        finalBuffer[totalSize - 1] = 0;

        return new BsonDocument(finalBuffer);
    }

    private void EnsureCapacity(int additional)
    {
        if (_position + additional > _buffer.Length)
        {
            var newBuffer = new byte[_buffer.Length * 2];
            _buffer.CopyTo(newBuffer, 0);
            _buffer = newBuffer;
        }
    }
}
