// Copyright 2020-2026 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0

namespace ZB.MOM.NatsNet.Server.Mqtt;

internal sealed class MqttWriter
{
    private readonly List<byte> _buffer;

    public MqttWriter(int capacity)
    {
        _buffer = new List<byte>(Math.Max(capacity, 0));
    }

    public byte[] Bytes()
    {
        return [.. _buffer];
    }

    public void WriteByte(byte value)
    {
        _buffer.Add(value);
    }

    public void Write(ReadOnlySpan<byte> value)
    {
        for (var i = 0; i < value.Length; i++)
        {
            _buffer.Add(value[i]);
        }
    }

    public void WriteUInt16(ushort value)
    {
        WriteByte((byte)(value >> 8));
        WriteByte((byte)value);
    }

    public void WriteString(string value)
    {
        WriteBytes(System.Text.Encoding.UTF8.GetBytes(value));
    }

    public void WriteBytes(ReadOnlySpan<byte> value)
    {
        WriteUInt16((ushort)value.Length);
        Write(value);
    }

    public void WriteVarInt(int value)
    {
        while (true)
        {
            var b = (byte)(value & 0x7F);
            value >>= 7;
            if (value > 0)
            {
                b |= 0x80;
            }

            WriteByte(b);
            if (value == 0)
            {
                break;
            }
        }
    }

    public byte WritePublishHeader(ushort packetIdentifier, byte qos, bool duplicate, bool retained, ReadOnlySpan<byte> topic, int payloadLength)
    {
        var packetLength = 2 + topic.Length + payloadLength;
        byte flags = 0;

        if (duplicate)
        {
            flags |= MqttPubFlag.Dup;
        }

        if (retained)
        {
            flags |= MqttPubFlag.Retain;
        }

        if (qos > 0)
        {
            packetLength += 2;
            flags |= (byte)(qos << 1);
        }

        WriteByte((byte)(MqttPacket.Pub | flags));
        WriteVarInt(packetLength);
        WriteBytes(topic);
        if (qos > 0)
        {
            WriteUInt16(packetIdentifier);
        }

        return flags;
    }

    public static (byte Flags, byte[] Header) MqttMakePublishHeader(ushort packetIdentifier, byte qos, bool duplicate, bool retained, ReadOnlySpan<byte> topic, int payloadLength)
    {
        var writer = NewMqttWriter(MqttConst.InitialPubHeader + topic.Length);
        var flags = writer.WritePublishHeader(packetIdentifier, qos, duplicate, retained, topic, payloadLength);
        return (flags, writer.Bytes());
    }

    public static MqttWriter NewMqttWriter(int capacity)
    {
        return new MqttWriter(capacity);
    }
}
