using System.Buffers;
using System.Buffers.Binary;
using Google.Protobuf;
using ZB.MOM.WW.MxGateway.Contracts.Proto;

namespace ZB.MOM.WW.MxGateway.Server.Workers;

/// <summary>
/// Writes length-prefixed WorkerEnvelope protobuf messages to a stream.
/// </summary>
public sealed class WorkerFrameWriter
{
    private readonly WorkerFrameProtocolOptions _options;
    private readonly Stream _stream;

    /// <summary>
    /// Initializes the writer with a stream and frame protocol options.
    /// </summary>
    /// <param name="stream">Stream to write frames to.</param>
    /// <param name="options">Frame protocol configuration.</param>
    public WorkerFrameWriter(
        Stream stream,
        WorkerFrameProtocolOptions options)
    {
        _stream = stream ?? throw new ArgumentNullException(nameof(stream));
        _options = options ?? throw new ArgumentNullException(nameof(options));
    }

    /// <summary>
    /// Writes a WorkerEnvelope as a length-prefixed message to the stream.
    /// </summary>
    /// <param name="envelope">Worker envelope message to write.</param>
    /// <param name="cancellationToken">Token to cancel the asynchronous operation.</param>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public async ValueTask WriteAsync(
        WorkerEnvelope envelope,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(envelope);
        WorkerEnvelopeValidator.Validate(envelope, _options);

        int payloadLength = envelope.CalculateSize();
        if (payloadLength == 0)
        {
            throw new WorkerFrameProtocolException(
                WorkerFrameProtocolErrorCode.InvalidEnvelope,
                "Worker envelope cannot serialize to an empty payload.");
        }

        if (payloadLength > _options.MaxMessageBytes)
        {
            throw new WorkerFrameProtocolException(
                WorkerFrameProtocolErrorCode.MessageTooLarge,
                $"Worker envelope payload length {payloadLength} exceeds the configured maximum of {_options.MaxMessageBytes} bytes.");
        }

        // Serialize once into a single pooled buffer that carries the 4-byte little-endian length
        // prefix followed by the payload, then issue one stream write. This avoids a second
        // serialization pass (ToByteArray re-runs CalculateSize), a separate prefix array and its
        // own write, and any per-frame heap allocation (GWC-08, IPC-13). The rented buffer may be
        // larger than requested, so only the first frameLength bytes are ever written.
        int frameLength = sizeof(uint) + payloadLength;
        byte[] frame = ArrayPool<byte>.Shared.Rent(frameLength);
        try
        {
            BinaryPrimitives.WriteUInt32LittleEndian(frame, (uint)payloadLength);
            envelope.WriteTo(new Span<byte>(frame, sizeof(uint), payloadLength));

            await _stream
                .WriteAsync(new ReadOnlyMemory<byte>(frame, 0, frameLength), cancellationToken)
                .ConfigureAwait(false);
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(frame);
        }
    }
}