// SPDX-License-Identifier: GPL-3.0-or-later

#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif

#include "libnetdata/libnetdata.h"
void pulse_aclk_sent_message_acked(usec_t publish_latency, size_t len);

#include "common_internal.h"
#include "mqtt_constants.h"
#include "mqtt_ng.h"
#include "aclk_mqtt_workers.h"
#include "daemon/config/netdata-conf-profile.h"

#define PACKET_ACK_TIMEOUT_SECS (60)
#define SMALL_STRING_DONT_FRAGMENT_LIMIT 128

#define LOCK_HDR_BUFFER(buffer) spinlock_lock(&((buffer)->spinlock))
#define UNLOCK_HDR_BUFFER(buffer) spinlock_unlock(&((buffer)->spinlock))

#define BUFFER_FRAG_GARBAGE_COLLECT         0x01
// some packets can be marked for garbage collection
// immediately when they are sent (e.g. sent PUBACK on QoS1)
#define BUFFER_FRAG_GARBAGE_COLLECT_ON_SEND 0x02
// as buffer fragment can point to both
// external data and data in the same buffer
// we mark the former case with BUFFER_FRAG_DATA_EXTERNAL
#define BUFFER_FRAG_DATA_EXTERNAL           0x04
// as single MQTT Packet can be stored into multiple
// buffer fragments (depending on copy requirements)
// this marks this fragment to be the first/last
#define BUFFER_FRAG_MQTT_PACKET_HEAD        0x10
#define BUFFER_FRAG_MQTT_PACKET_TAIL        0x20

typedef uint16_t buffer_frag_flag_t;
struct buffer_fragment {
    uint32_t len;
    uint32_t sent;
    buffer_frag_flag_t flags;
    uint16_t packet_id;
    void (*free_fnc)(void *ptr);
    unsigned char *data;
    // timestamp (monotonic usec) of when this MQTT packet was enqueued into the transaction buffer (set on HEAD)
    usec_t enqueued_monotonic_ut;
    usec_t sent_monotonic_ut;
    struct buffer_fragment *next;
};

typedef struct buffer_fragment *mqtt_msg_data;

// buffer used for MQTT headers only
// not for actual data sent
struct header_buffer {
    size_t size;
    unsigned char *data;
    unsigned char *tail;
    struct buffer_fragment *tail_frag;
};

struct transaction_buffer {
    struct header_buffer hdr_buffer;
    // used while building new message
    // to be able to revert state easily
    // in case of error mid processing
    struct header_buffer state_backup;
    SPINLOCK spinlock;
    struct buffer_fragment *sending_frag;
    // Fragment owned outside hdr_buffer that remains valid across compaction/realloc.
    struct buffer_fragment *stable_frag;
};

enum mqtt_client_state {
    MQTT_STATE_RAW = 0,
    MQTT_STATE_CONNECT_PENDING,
    MQTT_STATE_CONNECTING,
    MQTT_STATE_CONNECTED,
    MQTT_STATE_ERROR,
    MQTT_STATE_DISCONNECTED
};

enum parser_state {
    MQTT_PARSE_FIXED_HEADER_PACKET_TYPE = 0,
    MQTT_PARSE_FIXED_HEADER_LEN,
    MQTT_PARSE_VARIABLE_HEADER,
    MQTT_PARSE_MQTT_PACKET_DONE
};

enum varhdr_parser_state {
    MQTT_PARSE_VARHDR_INITIAL = 0,
    MQTT_PARSE_VARHDR_OPTIONAL_REASON_CODE,
    MQTT_PARSE_VARHDR_PROPS,
    MQTT_PARSE_VARHDR_TOPICNAME,
    MQTT_PARSE_VARHDR_POST_TOPICNAME,
    MQTT_PARSE_VARHDR_PACKET_ID,
    MQTT_PARSE_REASONCODES,
    MQTT_PARSE_PAYLOAD
};

struct mqtt_vbi_parser_ctx {
    char data[MQTT_VBI_MAXBYTES];
    uint8_t bytes;
    uint32_t result;
};

enum mqtt_datatype {
    MQTT_TYPE_UNKNOWN = 0,
    MQTT_TYPE_UINT_8,
    MQTT_TYPE_UINT_16,
    MQTT_TYPE_UINT_32,
    MQTT_TYPE_VBI,
    MQTT_TYPE_STR,
    MQTT_TYPE_STR_PAIR,
    MQTT_TYPE_BIN
};

struct mqtt_property {
    uint8_t id;
    enum mqtt_datatype type;
    union {
        char *strings[2];
        void *bindata;
        uint8_t uint8;
        uint16_t uint16;
        uint32_t uint32;
    } data;
    size_t bindata_len;
    struct mqtt_property *next;
};

enum mqtt_properties_parser_state {
    PROPERTIES_LENGTH = 0,
    PROPERTY_CREATE,
    PROPERTY_ID,
    PROPERTY_TYPE_UINT8,
    PROPERTY_TYPE_UINT16,
    PROPERTY_TYPE_UINT32,
    PROPERTY_TYPE_STR_BIN_LEN,
    PROPERTY_TYPE_STR,
    PROPERTY_TYPE_BIN,
    PROPERTY_TYPE_VBI,
    PROPERTY_NEXT
};

struct mqtt_properties_parser_ctx {
    enum mqtt_properties_parser_state state;
    struct mqtt_property *head;
    struct mqtt_property *tail;
    uint32_t properties_length;
    uint32_t vbi_length;
    size_t max_bytes;
    struct mqtt_vbi_parser_ctx vbi_parser_ctx;
    size_t bytes_consumed;
    int str_idx;
};

struct mqtt_connack {
    uint8_t flags;
    uint8_t reason_code;
};
struct mqtt_puback {
    uint16_t packet_id;
    uint8_t reason_code;
};

struct mqtt_suback {
    uint16_t packet_id;
    uint8_t *reason_codes;
    uint8_t reason_code_count;
    uint8_t reason_codes_pending;
};

struct mqtt_publish {
    uint16_t topic_len;
    char *topic;
    uint16_t packet_id;
    size_t data_len;
    char *data;
    uint8_t qos;
};

struct mqtt_disconnect {
    uint8_t reason_code;
};

struct mqtt_ng_parser {
    rbuf_t received_data;

    uint8_t mqtt_control_packet_type;
    uint32_t mqtt_fixed_hdr_remaining_length;
    size_t mqtt_parsed_len;

    struct mqtt_vbi_parser_ctx vbi_parser;
    struct mqtt_properties_parser_ctx properties_parser;

    enum parser_state state;
    enum varhdr_parser_state varhdr_state;

    struct mqtt_property *varhdr_properties;

    union {
        struct mqtt_connack connack;
        struct mqtt_puback puback;
        struct mqtt_suback suback;
        struct mqtt_publish publish;
        struct mqtt_disconnect disconnect;
    } mqtt_packet;
};

struct topic_alias_data {
    uint16_t idx;
    uint32_t usage_count;
};

struct topic_aliases_data {
    c_rhash stoi_dict;
    uint32_t idx_max;
    uint32_t idx_assigned;
    SPINLOCK spinlock;
};

struct mqtt_ng_client {
    struct transaction_buffer main_buffer;

    enum mqtt_client_state client_state;

    mqtt_msg_data connect_msg;

    mqtt_ng_send_fnc_t send_fnc_ptr;
    void *user_ctx;

    // time when last fragment of MQTT message was sent
    time_t time_of_last_send;

    struct mqtt_ng_parser parser;

    size_t max_mem_bytes;

    void (*puback_callback)(uint16_t packet_id);
    void (*connack_callback)(void* user_ctx, int connack_reply);
    void (*msg_callback)(const char *topic, const void *msg, size_t msglen, int qos);

    unsigned int ping_pending:1;
    struct buffer_fragment ping_frag;

    struct mqtt_ng_stats stats;

    struct {
        SPINLOCK spinlock;
        Pvoid_t JudyL;
    } pending_packets;

    struct topic_aliases_data tx_topic_aliases;
    c_rhash rx_aliases;

    size_t max_msg_size;

    // MQTT 5.0 server Receive Maximum (CONNACK prop 0x21); absent => 65535 default [MQTT-3.2.2.3.3]
    uint16_t rx_maximum;
};

usec_t publish_latency;

static unsigned char pingreq[] = { MQTT_CPT_PINGREQ << 4, 0x00 };

int uint32_to_mqtt_vbi(uint32_t input, unsigned char *output) {
    int i = 1;
    *output = 0;

    /* MQTT 5 specs allows max 4 bytes of output
       making it 0xFF, 0xFF, 0xFF, 0x7F
       representing number 268435455 decimal
       see 1.5.5. Variable Byte Integer */
    if(input >= 256 * 1024 * 1024)
        return 0;

    if(!input) {
        *output = 0;
        return 1;
    }

    while(input) {
        output[i-1] = input & MQTT_VBI_DATA_MASK;
        input >>= 7;
        if (input)
            output[i-1] |= MQTT_VBI_CONTINUATION_FLAG;
        i++;
    }
    return i - 1;
}

int mqtt_vbi_to_uint32(char *input, uint32_t *output) {
    // dont want to operate directly on output
    // as I want it to be possible for input and output
    // pointer to be the same
    uint32_t result = 0;
    uint32_t multiplier = 1;

    do {
        result += (uint32_t)(*input & MQTT_VBI_DATA_MASK) * multiplier;
        if (multiplier > 128*128*128)
            return 1;
        multiplier <<= 7;
    } while (*input++ & MQTT_VBI_CONTINUATION_FLAG);
    *output = result;
    return 0;
}

#ifdef TESTS
#include <stdio.h>
#define MQTT_VBI_MAXLEN 4
// we add extra byte to check we dont write out of bounds
// in case where 4 bytes are supposed to be written
static const char _mqtt_vbi_0[MQTT_VBI_MAXLEN + 1] = { 0x00, 0x00, 0x00, 0x00, 0x00 };
static const char _mqtt_vbi_127[MQTT_VBI_MAXLEN + 1] = { 0x7F, 0x00, 0x00, 0x00, 0x00 };
static const char _mqtt_vbi_128[MQTT_VBI_MAXLEN + 1] = { 0x80, 0x01, 0x00, 0x00, 0x00 };
static const char _mqtt_vbi_16383[MQTT_VBI_MAXLEN + 1] = { 0xFF, 0x7F, 0x00, 0x00, 0x00 };
static const char _mqtt_vbi_16384[MQTT_VBI_MAXLEN + 1] = { 0x80, 0x80, 0x01, 0x00, 0x00 };
static const char _mqtt_vbi_2097151[MQTT_VBI_MAXLEN + 1] = { 0xFF, 0xFF, 0x7F, 0x00, 0x00 };
static const char _mqtt_vbi_2097152[MQTT_VBI_MAXLEN + 1] = { 0x80, 0x80, 0x80, 0x01, 0x00 };
static const char _mqtt_vbi_268435455[MQTT_VBI_MAXLEN + 1] = { 0xFF, 0xFF, 0xFF, 0x7F, 0x00 };
static const char _mqtt_vbi_999999999[MQTT_VBI_MAXLEN + 1] = { 0x80, 0x80, 0x80, 0x80, 0x01 };

#define MQTT_VBI_TESTCASE(case, expected_len) \
    { \
    memset(buf, 0, MQTT_VBI_MAXLEN + 1); \
    int len; \
    if ((len=uint32_to_mqtt_vbi(case, buf)) != expected_len) { \
        fprintf(stderr, "uint32_to_mqtt_vbi(case:%d, line:%d): Incorrect length returned. Expected %d, Got %d\n", case, __LINE__, expected_len, len); \
        return 1; \
    } \
    if (memcmp(buf, _mqtt_vbi_ ## case, MQTT_VBI_MAXLEN + 1 )) { \
        fprintf(stderr, "uint32_to_mqtt_vbi(case:%d, line:%d): Wrong output\n", case, __LINE__); \
        return 1; \
    } }


int test_uint32_mqtt_vbi() {
    char buf[MQTT_VBI_MAXLEN + 1];

    MQTT_VBI_TESTCASE(0,         1)
    MQTT_VBI_TESTCASE(127,       1)
    MQTT_VBI_TESTCASE(128,       2)
    MQTT_VBI_TESTCASE(16383,     2)
    MQTT_VBI_TESTCASE(16384,     3)
    MQTT_VBI_TESTCASE(2097151,   3)
    MQTT_VBI_TESTCASE(2097152,   4)
    MQTT_VBI_TESTCASE(268435455, 4)

    memset(buf, 0, MQTT_VBI_MAXLEN + 1);
    int len;
    if ((len=uint32_to_mqtt_vbi(268435456, buf)) != 0) {
        fprintf(stderr, "uint32_to_mqtt_vbi(case:268435456, line:%d): Incorrect length returned. Expected 0, Got %d\n", __LINE__, len);
        return 1;
    }

    return 0;
}

#define MQTT_VBI2UINT_TESTCASE(case, expected_error) \
    { \
    uint32_t result; \
    int ret = mqtt_vbi_to_uint32(_mqtt_vbi_ ## case, &result); \
    if (ret && !(expected_error)) { \
        fprintf(stderr, "mqtt_vbi_to_uint(case:%d, line:%d): Unexpectedly Errored\n", (case), __LINE__); \
        return 1; \
    } \
    if (!ret && (expected_error)) { \
        fprintf(stderr, "mqtt_vbi_to_uint(case:%d, line:%d): Should return error but didnt\n", (case), __LINE__); \
        return 1; \
    } \
    if (!ret && result != (case)) { \
        fprintf(stderr, "mqtt_vbi_to_uint(case:%d, line:%d): Returned wrong result %d\n", (case), __LINE__, result); \
        return 1; \
    }}


int test_mqtt_vbi_to_uint32() {
    MQTT_VBI2UINT_TESTCASE(0,         0)
    MQTT_VBI2UINT_TESTCASE(127,       0)
    MQTT_VBI2UINT_TESTCASE(128,       0)
    MQTT_VBI2UINT_TESTCASE(16383,     0)
    MQTT_VBI2UINT_TESTCASE(16384,     0)
    MQTT_VBI2UINT_TESTCASE(2097151,   0)
    MQTT_VBI2UINT_TESTCASE(2097152,   0)
    MQTT_VBI2UINT_TESTCASE(268435455, 0)
    MQTT_VBI2UINT_TESTCASE(999999999, 1)
    return 0;
}
#endif /* TESTS */

// this helps with switch statements
// as they have to use integer type (not pointer)
enum memory_mode {
    MEMCPY,
    EXTERNAL_FREE_AFTER_USE,
    CALLER_RESPONSIBLE
};

static enum memory_mode ptr2memory_mode(void * ptr) {
    if (ptr == NULL)
        return MEMCPY;
    if (ptr == CALLER_RESPONSIBILITY)
        return CALLER_RESPONSIBLE;
    return EXTERNAL_FREE_AFTER_USE;
}

#define frag_is_marked_for_gc(frag) ((frag->flags & BUFFER_FRAG_GARBAGE_COLLECT) || ((frag->flags & BUFFER_FRAG_GARBAGE_COLLECT_ON_SEND) && frag->sent == frag->len))
#define FRAG_SIZE_IN_BUFFER(frag) (sizeof(struct buffer_fragment) + ((frag->flags & BUFFER_FRAG_DATA_EXTERNAL) ? 0 : frag->len))

static void buffer_frag_free_data(struct buffer_fragment *frag)
{
    if ( frag->flags & BUFFER_FRAG_DATA_EXTERNAL && frag->data != NULL) {
        switch (ptr2memory_mode(frag->free_fnc)) {
            case MEMCPY:
                freez(frag->data);
                break;
            case EXTERNAL_FREE_AFTER_USE:
                frag->free_fnc(frag->data);
                break;
            case CALLER_RESPONSIBLE:
                break;
        }
        frag->data = NULL;
    }
}

#define HEADER_BUFFER_SIZE (1024*1024)
#define HEADER_BUFFER_SIZE_IOT (128*1024)
#define HEADER_BUFFER_SIZE_STANDALONE (512*1024)

#define GROWTH_FACTOR 1.25

#define BUFFER_BYTES_USED(buf) ((size_t)((buf)->tail - (buf)->data))
#define BUFFER_BYTES_AVAILABLE(buf) ((buf)->size - BUFFER_BYTES_USED(buf))
#define BUFFER_FIRST_FRAG(buf) ((struct buffer_fragment *)((buf)->tail_frag ? (buf)->data : NULL))
static void buffer_purge(struct header_buffer *buf) {
    struct buffer_fragment *frag = BUFFER_FIRST_FRAG(buf);
    while (frag) {
        buffer_frag_free_data(frag);
        frag = frag->next;
    }
    buf->tail = buf->data;
    buf->tail_frag = NULL;
}

#define FRAG_PADDING(addr) ((MQTT_WSS_FRAG_MEMALIGN - ((uintptr_t)addr % MQTT_WSS_FRAG_MEMALIGN)) % MQTT_WSS_FRAG_MEMALIGN)
static struct buffer_fragment *buffer_new_frag(struct header_buffer *buf, buffer_frag_flag_t flags)
{
    uint8_t padding = FRAG_PADDING(buf->tail);

    if (BUFFER_BYTES_AVAILABLE(buf) < sizeof(struct buffer_fragment) + padding)
        return NULL;

    struct buffer_fragment *frag = (struct buffer_fragment *)(buf->tail + padding);

    memset(frag, 0, sizeof(*frag));
    buf->tail += sizeof(*frag) + padding;

    if (/*!((frag)->flags & BUFFER_FRAG_MQTT_PACKET_HEAD) &&*/ buf->tail_frag)
        buf->tail_frag->next = frag;

    buf->tail_frag = frag;

    frag->data = buf->tail;

    frag->flags = flags;

    return frag;
}

static void buffer_rebuild(struct header_buffer *buf)
{
    struct buffer_fragment *frag = (struct buffer_fragment*)buf->data;
    do {
        buf->tail = (unsigned char *) frag + sizeof(struct buffer_fragment);
        buf->tail_frag = frag;
        if (!(frag->flags & BUFFER_FRAG_DATA_EXTERNAL)) {
            buf->tail_frag->data = buf->tail;
            buf->tail += frag->len;
        }
        if (frag->next != NULL)
            frag->next = (struct buffer_fragment*)(buf->tail + FRAG_PADDING(buf->tail));
        frag = frag->next;
    } while(frag);
}

static void buffer_garbage_collect(struct header_buffer *buf, bool main_thread)
{
    struct buffer_fragment *frag = BUFFER_FIRST_FRAG(buf);
    while (frag) {
        if (!frag_is_marked_for_gc(frag))
            break;

        buffer_frag_free_data(frag);

        frag = frag->next;
    }

    if (frag == BUFFER_FIRST_FRAG(buf))
        return;

    if (!frag) {
        buf->tail_frag = NULL;
        buf->tail = buf->data;
        return;
    }

#ifdef ADDITIONAL_CHECKS
    if (!(frag->flags & BUFFER_FRAG_MQTT_PACKET_HEAD)) {
        mws_error(log_ctx, "Expected to find end of buffer (NULL) or next packet head!");
        return;
    }
#endif

    memmove(buf->data, frag, buf->tail - (unsigned char *) frag);
    if (main_thread)
        worker_is_busy(WORKER_ACLK_BUFFER_COMPACT);

    buffer_rebuild(buf);

    if (main_thread)
        worker_is_idle();
}

static void transaction_buffer_garbage_collect(struct transaction_buffer *buf, bool main_thread)
{
    if (main_thread)
        worker_is_busy(WORKER_ACLK_RECLAIM_MEMORY);
    // Invalidate the cached sending fragment
    // as we will move data around
    if (buf->sending_frag != buf->stable_frag)
        buf->sending_frag = NULL;

    buffer_garbage_collect(&buf->hdr_buffer, main_thread);
    if (main_thread)
        worker_is_idle();
}

static int transaction_buffer_grow(struct transaction_buffer *buf, float rate, size_t max)
{
    if (buf->hdr_buffer.size >= max)
        return 0;

    // Invalidate the cached sending fragment
    // as we will move data around
    if (buf->sending_frag != buf->stable_frag)
        buf->sending_frag = NULL;

    buf->hdr_buffer.size = (size_t)((float)buf->hdr_buffer.size * rate);
    if (buf->hdr_buffer.size > max)
        buf->hdr_buffer.size = max;

    void *ret = reallocz(buf->hdr_buffer.data, buf->hdr_buffer.size);
    if (ret == NULL) {
        nd_log(NDLS_DAEMON, NDLP_WARNING, "Buffer growth failed (realloc)");
        return 1;
    }

    nd_log(NDLS_DAEMON, NDLP_DEBUG, "Message metadata buffer was grown");

    buf->hdr_buffer.data = ret;
    buffer_rebuild(&buf->hdr_buffer);
    return 0;
}

inline static void transaction_buffer_init(struct transaction_buffer *to_init, size_t size)
{
    spinlock_init(&to_init->spinlock);

    to_init->hdr_buffer.size = size;
    to_init->hdr_buffer.data = mallocz(size);
    to_init->hdr_buffer.tail = to_init->hdr_buffer.data;
    to_init->hdr_buffer.tail_frag = NULL;
    to_init->stable_frag = NULL;
}

static void transaction_buffer_destroy(struct transaction_buffer *to_init)
{
    buffer_purge(&to_init->hdr_buffer);
    freez(to_init->hdr_buffer.data);
}

static void mqtt_ng_init_ping_fragment(struct mqtt_ng_client *client)
{
    client->ping_frag = (struct buffer_fragment){
        .data = pingreq,
        .flags = BUFFER_FRAG_MQTT_PACKET_HEAD | BUFFER_FRAG_MQTT_PACKET_TAIL,
        .free_fnc = NULL,
        .len = sizeof(pingreq),
        .next = NULL,
        .sent = 0,
        .packet_id = 0,
    };
    client->main_buffer.stable_frag = &client->ping_frag;
}

// Creates transaction
// saves state of buffer before any operation was done
// allowing for rollback if things go wrong
#define transaction_buffer_transaction_start(buf)                                                                      \
    {                                                                                                                  \
        LOCK_HDR_BUFFER(buf);                                                                                          \
        memcpy(&(buf)->state_backup, &(buf)->hdr_buffer, sizeof((buf)->hdr_buffer));                                   \
    }

#define transaction_buffer_transaction_commit(buf) UNLOCK_HDR_BUFFER(buf);

void transaction_buffer_transaction_rollback(struct transaction_buffer *buf, struct buffer_fragment *frag)
{
    memcpy(&buf->hdr_buffer, &buf->state_backup, sizeof(buf->hdr_buffer));
    if (buf->hdr_buffer.tail_frag != NULL)
        buf->hdr_buffer.tail_frag->next = NULL;

    while(frag) {
        buffer_frag_free_data(frag);
        // we are not actually freeing the structure itself
        // just the data it manages
        // structure itself is in permanent buffer
        // which is locked by HDR_BUFFER lock
        frag = frag->next;
    }

    UNLOCK_HDR_BUFFER(buf);
}

#define TX_ALIASES_INITIALIZE() c_rhash_new(0)
#define RX_ALIASES_INITIALIZE() c_rhash_new(UINT16_MAX >> 8)
struct mqtt_ng_client *mqtt_ng_init(struct mqtt_ng_init *settings)
{
    struct mqtt_ng_client *client = callocz(1, sizeof(struct mqtt_ng_client));

    size_t buffer_size = netdata_conf_is_iot() ?
                             HEADER_BUFFER_SIZE_IOT :
                             (netdata_conf_is_standalone() ? HEADER_BUFFER_SIZE_STANDALONE : HEADER_BUFFER_SIZE);
    transaction_buffer_init(&client->main_buffer, buffer_size);
    mqtt_ng_init_ping_fragment(client);

    client->rx_aliases = RX_ALIASES_INITIALIZE();

    spinlock_init(&client->tx_topic_aliases.spinlock);

    client->tx_topic_aliases.stoi_dict = TX_ALIASES_INITIALIZE();
    client->tx_topic_aliases.idx_max = UINT16_MAX;

    // MQTT 5.0 default Receive Maximum when the server omits the property [MQTT-3.2.2.3.3]
    __atomic_store_n(&client->rx_maximum, UINT16_MAX, __ATOMIC_RELAXED);

    // TODO just embed the struct into mqtt_ng_client
    client->parser.received_data = settings->data_in;
    client->send_fnc_ptr = settings->data_out_fnc;
    client->user_ctx = settings->user_ctx;

    client->puback_callback = settings->puback_callback;
    client->connack_callback = settings->connack_callback;
    client->msg_callback = settings->msg_callback;
    spinlock_init(&client->pending_packets.spinlock);
    client->pending_packets.JudyL = NULL;
    __atomic_store_n(&publish_latency, 0, __ATOMIC_RELEASE);

    return client;
}

static uint8_t get_control_packet_type(uint8_t first_hdr_byte)
{
    return first_hdr_byte >> 4;
}

static void mqtt_ng_parser_reset(struct mqtt_ng_parser *parser);

static void mqtt_ng_destroy_rx_alias_hash(c_rhash hash)
{
    c_rhash_iter_t i = C_RHASH_ITER_T_INITIALIZER;
    uint64_t stored_key;
    void *to_free;
    while(!c_rhash_iter_uint64_keys(hash, &i, &stored_key)) {
        c_rhash_get_ptr_by_uint64(hash, stored_key, &to_free);
        freez(to_free);
    }
    c_rhash_destroy(hash);
}

static void mqtt_ng_destroy_tx_alias_hash(c_rhash hash)
{
    c_rhash_iter_t i = C_RHASH_ITER_T_INITIALIZER;
    const char *stored_key;
    void *to_free;
    while(!c_rhash_iter_str_keys(hash, &i, &stored_key)) {
        c_rhash_get_ptr_by_str(hash, stored_key, &to_free);
        freez(to_free);
    }
    c_rhash_destroy(hash);
}

static void destroy_timeout_monitor_list(struct mqtt_ng_client *client)
{
    spinlock_lock(&client->pending_packets.spinlock);
    (void) JudyLFreeArray(&client->pending_packets.JudyL, PJE0);
    spinlock_unlock(&client->pending_packets.spinlock);
    __atomic_store_n(&client->stats.packets_waiting_puback, 0, __ATOMIC_RELAXED);
}

void mqtt_ng_destroy(struct mqtt_ng_client *client)
{
    mqtt_ng_parser_reset(&client->parser);
    transaction_buffer_destroy(&client->main_buffer);

    mqtt_ng_destroy_tx_alias_hash(client->tx_topic_aliases.stoi_dict);
    mqtt_ng_destroy_rx_alias_hash(client->rx_aliases);
    destroy_timeout_monitor_list(client);
    freez(client);
}

int frag_set_external_data(struct buffer_fragment *frag, void *data, size_t data_len, free_fnc_t data_free_fnc)
{
    if (frag->len) {
        // TODO?: This could potentially be done in future if we set rule
        // external data always follows in buffer data
        // could help reduce fragmentation in some messages but
        // currently not worth it considering time is tight
        nd_log(NDLS_DAEMON, NDLP_ERR, "INTERNAL ERROR: Cannot set external data to fragment already containing in buffer data!");
        return 1;
    }

    switch (ptr2memory_mode(data_free_fnc)) {
        case MEMCPY:
            frag->data = mallocz(data_len);
            memcpy(frag->data, data, data_len);
            break;
        case EXTERNAL_FREE_AFTER_USE:
        case CALLER_RESPONSIBLE:
            frag->data = data;
            break;
    }
    frag->free_fnc = data_free_fnc;
    frag->len = data_len;

    frag->flags |= BUFFER_FRAG_DATA_EXTERNAL;
    return 0;
 }

// this is fixed part of variable header for connect packet
// mqtt-v5.0-cs1, 3.1.2.1, 2.1.2.2
static const char mqtt_protocol_name_frag[] =
    { 0x00, 0x04, 'M', 'Q', 'T', 'T', MQTT_VERSION_5_0 };

// see 1.5.5
#define MQTT_VARSIZE_INT_BYTES(value) ( value > 2097152 ? 4 : ( value > 16384 ? 3 : ( value > 128 ? 2 : 1 ) ) )

static size_t mqtt_ng_connect_size(struct mqtt_auth_properties *auth,
                    struct mqtt_lwt_properties *lwt,
                    uint16_t client_id_len,
                    uint16_t will_topic_len,
                    uint16_t username_len,
                    uint16_t password_len)
{
    // First get the size of payload + variable header
    size_t size =
        + sizeof(mqtt_protocol_name_frag) /* Proto Name and Version */
        + 1 /* Connect Flags */
        + 2 /* Keep Alive */
        + 4 /* 3.1.2.11.1 Property Length - for now fixed to only Topic Alias Maximum, TODO TODO*/;

    // CONNECT payload. 3.1.3
    if (auth->client_id)
        size += 2 + client_id_len;

    if (lwt) {
        // 3.1.3.2 will properties TODO TODO
        size += 1;

        // 3.1.3.3
        if (lwt->will_topic)
            size += 2 + will_topic_len;

        // 3.1.3.4 will payload
        if (lwt->will_message) {
            size += 2 + lwt->will_message_size;
        }
    }

    // 3.1.3.5
    if (auth->username)
        size += 2 + username_len;

    // 3.1.3.6
    if (auth->password)
        size += 2 + password_len;

    return size;
}

#define BUFFER_TRANSACTION_NEW_FRAG(buf, flags, frag, on_fail) \
    { if(frag==NULL) { \
        frag = buffer_new_frag(buf, (flags)); } \
      if(frag==NULL) { on_fail; }}

#define CHECK_BYTES_AVAILABLE(buf, needed, fail)                                                                       \
    {                                                                                                                  \
        if (BUFFER_BYTES_AVAILABLE(buf) < (size_t)needed) {                                                            \
            fail;                                                                                                      \
        }                                                                                                              \
    }

#define DATA_ADVANCE(buf, bytes, frag) { size_t b = (bytes); (buf)->tail += b; (frag)->len += b; }

// TODO maybe just user client->buf.tail?
#define WRITE_POS(frag) (&(frag->data[frag->len]))

// [MQTT-1.5.2] Two Byte Integer
#define PACK_2B_INT(buffer, integer, frag)                                                                             \
    {                                                                                                                  \
        uint16_t temp = htobe16((integer));                                                                            \
        memcpy(WRITE_POS(frag), &temp, sizeof(uint16_t));                                                              \
        DATA_ADVANCE(buffer, sizeof(uint16_t), frag);                                                                  \
    }

static bool mqtt_ng_get_2byte_field_length(const char *field, const char *value, uint16_t *length)
{
    size_t len = strnlen(value, (size_t)UINT16_MAX + 1);
    if (unlikely(len > UINT16_MAX)) {
        nd_log(NDLS_DAEMON, NDLP_ERR, "MQTT %s exceeds the 65535-byte protocol limit", field);
        return false;
    }

    *length = (uint16_t)len;
    return true;
}

static int optimized_add(struct header_buffer *buf, void *data, size_t data_len, free_fnc_t data_free_fnc, struct buffer_fragment **frag)
{
    if (data_len > SMALL_STRING_DONT_FRAGMENT_LIMIT) {
        buffer_frag_flag_t flags = BUFFER_FRAG_DATA_EXTERNAL;
        if ((*frag)->flags & BUFFER_FRAG_GARBAGE_COLLECT_ON_SEND)
            flags |= BUFFER_FRAG_GARBAGE_COLLECT_ON_SEND;
        if( (*frag = buffer_new_frag(buf, flags)) == NULL ) {
           nd_log(NDLS_DAEMON, NDLP_ERR, "Out of buffer space while generating the message");
            return 1;
        }
        if (frag_set_external_data(*frag, data, data_len, data_free_fnc)) {
           nd_log(NDLS_DAEMON, NDLP_ERR, "Error adding external data to newly created fragment");
            return 1;
        }
        // we dont want to write to this fragment anymore
        *frag = NULL;
    } else if (data_len) {
        // if the data are small dont bother creating new fragments
        // store in buffer directly
        CHECK_BYTES_AVAILABLE(buf, data_len, return 1)
        memcpy(buf->tail, data, data_len);
        DATA_ADVANCE(buf, data_len, *frag)
    }
    return 0;
}

static void remove_packet_from_timeout_monitor_list_unsafe(struct mqtt_ng_client *client, uint16_t packet_id)
{
    int rc = JudyLDel(&client->pending_packets.JudyL, (Word_t) packet_id, PJE0);
    // rc = 1 if the packet was deleted, so update statistics
    if (likely(rc))
        __atomic_fetch_sub(&client->stats.packets_waiting_puback, 1, __ATOMIC_RELAXED);
}

#define PACKET_TIMEOUT_EPOCH (1704067200L)  // Jan 1, 2024 00:00:00 UTC

static void add_packet_to_timeout_monitor_list(struct mqtt_ng_client *client, uint16_t packet_id)
{
    spinlock_lock(&client->pending_packets.spinlock);
    time_t now = now_realtime_sec();
    // Add it to the JudyL array
    uint32_t *Pvalue = (uint32_t *) JudyLIns(&client->pending_packets.JudyL, (Word_t) packet_id, PJE0);
    if (Pvalue == PJERR) {
        nd_log(NDLS_DAEMON, NDLP_ERR, "Error inserting packet_id (%" PRIu16 ") into JudyL array.", packet_id);
        spinlock_unlock(&client->pending_packets.spinlock);
        return;
    }
    *Pvalue = (uint32_t) ((now - PACKET_TIMEOUT_EPOCH) + PACKET_ACK_TIMEOUT_SECS);
    spinlock_unlock(&client->pending_packets.spinlock);

    __atomic_fetch_add(&client->stats.packets_waiting_puback, 1, __ATOMIC_RELAXED);
}

#define TRY_GENERATE_MESSAGE(generator_function, ...)                                                                  \
    ({                                                                                                                 \
        int _rc = generator_function(&client->main_buffer, ##__VA_ARGS__);                                             \
        if (_rc == MQTT_NG_MSGGEN_BUFFER_OOM) {                                                                        \
            LOCK_HDR_BUFFER(&client->main_buffer);                                                                     \
            transaction_buffer_garbage_collect(&client->main_buffer, false);                                           \
            UNLOCK_HDR_BUFFER(&client->main_buffer);                                                                   \
            _rc = generator_function(&client->main_buffer, ##__VA_ARGS__);                                             \
            if (_rc == MQTT_NG_MSGGEN_BUFFER_OOM && client->max_mem_bytes) {                                           \
                LOCK_HDR_BUFFER(&client->main_buffer);                                                                 \
                transaction_buffer_grow((&client->main_buffer), GROWTH_FACTOR, client->max_mem_bytes);                 \
                UNLOCK_HDR_BUFFER(&client->main_buffer);                                                               \
                _rc = generator_function(&client->main_buffer, ##__VA_ARGS__);                                         \
            }                                                                                                          \
            if (_rc == MQTT_NG_MSGGEN_BUFFER_OOM)                                                                      \
                nd_log(                                                                                                \
                    NDLS_DAEMON,                                                                                       \
                    NDLP_ERR,                                                                                          \
                    "%s failed to generate message due to insufficient buffer space (line %d)",                        \
                    __FUNCTION__,                                                                                      \
                    __LINE__);                                                                                         \
        }                                                                                                              \
        if (_rc == MQTT_NG_MSGGEN_OK) {                                                                                \
            __atomic_fetch_add(&client->stats.tx_messages_queued, 1, __ATOMIC_RELAXED);                                \
        }                                                                                                              \
        _rc;                                                                                                           \
    })

mqtt_msg_data mqtt_ng_generate_connect(
    struct transaction_buffer *trx_buf,
    struct mqtt_auth_properties *auth,
    struct mqtt_lwt_properties *lwt,
    uint16_t keep_alive)
{
    // Sanity Checks First (are given parameters correct and up to MQTT spec)
    if (!auth->client_id) {
        nd_log(NDLS_DAEMON, NDLP_ERR, "ClientID must be set. [MQTT-3.1.3-3]");
        return NULL;
    }

    uint16_t client_id_len;
    if (!mqtt_ng_get_2byte_field_length("Client ID", auth->client_id, &client_id_len))
        return NULL;

    if (!client_id_len) {
        // [MQTT-3.1.3-6] server MAY allow empty client_id and treat it
        // as specific client_id (not same as client_id not given)
        // however server MUST allow ClientIDs between 1-23 bytes [MQTT-3.1.3-5]
        // so we will warn client server might not like this and he is using it
        // at his own risk!
        nd_log(NDLS_DAEMON, NDLP_WARNING, "client_id provided is empty string. This might not be allowed by server [MQTT-3.1.3-6]");
    }

    uint16_t will_topic_len = 0;

    if (lwt) {
        if (lwt->will_message && lwt->will_message_size > 65535) {
            nd_log(NDLS_DAEMON, NDLP_ERR, "Will message cannot be longer than 65535 bytes due to MQTT protocol limitations [MQTT-3.1.3-4] and [MQTT-1.5.6]");
            return NULL;
        }

        if (!lwt->will_topic) { //TODO topic given with strlen==0 ? check specs
            nd_log(NDLS_DAEMON, NDLP_ERR, "If will message is given will topic must also be given [MQTT-3.1.3.3]");
            return NULL;
        }

        if (!mqtt_ng_get_2byte_field_length("Will Topic", lwt->will_topic, &will_topic_len))
            return NULL;

        if (lwt->will_qos > MQTT_MAX_QOS) {
            // refer to [MQTT-3-1.2-12]
            nd_log(NDLS_DAEMON, NDLP_ERR, "QOS for LWT message is bigger than max");
            return NULL;
        }
    }

    uint16_t username_len = 0;
    if (auth->username && !mqtt_ng_get_2byte_field_length("User Name", auth->username, &username_len))
        return NULL;

    uint16_t password_len = 0;
    if (auth->password && !mqtt_ng_get_2byte_field_length("Password", auth->password, &password_len))
        return NULL;

    // >> START THE RODEO <<
    transaction_buffer_transaction_start(trx_buf)

    // Calculate the resulting message size sans fixed MQTT header
    size_t size = mqtt_ng_connect_size(
        auth, lwt, client_id_len, will_topic_len, username_len, password_len);

    // Start generating the message
    struct buffer_fragment *frag = NULL;
    mqtt_msg_data ret = NULL;

    BUFFER_TRANSACTION_NEW_FRAG(&trx_buf->hdr_buffer, BUFFER_FRAG_MQTT_PACKET_HEAD, frag, goto fail_rollback)
    ret = frag;

    // MQTT Fixed Header
    size_t needed_bytes = 1 /* Packet type */ + MQTT_VARSIZE_INT_BYTES(size) + sizeof(mqtt_protocol_name_frag) + 1 /* CONNECT FLAGS */ + 2 /* keepalive */ + 4 /* Properties TODO now fixed to Topic Alias Maximum */;
    CHECK_BYTES_AVAILABLE(&trx_buf->hdr_buffer, needed_bytes, goto fail_rollback)

    *WRITE_POS(frag) = MQTT_CPT_CONNECT << 4;
    DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)
    DATA_ADVANCE(&trx_buf->hdr_buffer, uint32_to_mqtt_vbi(size, WRITE_POS(frag)), frag)

    memcpy(WRITE_POS(frag), mqtt_protocol_name_frag, sizeof(mqtt_protocol_name_frag));
    DATA_ADVANCE(&trx_buf->hdr_buffer, sizeof(mqtt_protocol_name_frag), frag)

    // [MQTT-3.1.2.3] Connect flags
    unsigned char *connect_flags = WRITE_POS(frag);
    *connect_flags = 0;
    if (auth->username)
        *connect_flags |= MQTT_CONNECT_FLAG_USERNAME;

    if (auth->password)
        *connect_flags |= MQTT_CONNECT_FLAG_PASSWORD;

    if (lwt) {
        *connect_flags |= MQTT_CONNECT_FLAG_LWT;
        *connect_flags |= lwt->will_qos << MQTT_CONNECT_FLAG_QOS_BITSHIFT;
        if (lwt->will_retain)
            *connect_flags |= MQTT_CONNECT_FLAG_LWT_RETAIN;
    }

    *connect_flags |= MQTT_CONNECT_FLAG_CLEAN_START;

    DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)

    PACK_2B_INT(&trx_buf->hdr_buffer, keep_alive, frag)

    // TODO Property Length [MQTT-3.1.3.2.1] temporary fixed to 3 (one property topic alias max)
    DATA_ADVANCE(&trx_buf->hdr_buffer, uint32_to_mqtt_vbi(3, WRITE_POS(frag)), frag)
    *WRITE_POS(frag) = MQTT_PROP_TOPIC_ALIAS_MAX;
    DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)

    PACK_2B_INT(&trx_buf->hdr_buffer, 65535, frag)

    // [MQTT-3.1.3.1] Client identifier
    CHECK_BYTES_AVAILABLE(&trx_buf->hdr_buffer, 2, goto fail_rollback)
    PACK_2B_INT(&trx_buf->hdr_buffer, client_id_len, frag)
    if (optimized_add(&trx_buf->hdr_buffer, auth->client_id, client_id_len, auth->client_id_free, &frag))
        goto fail_rollback;

    if (lwt != NULL) {
        // Will Properties [MQTT-3.1.3.2]
        // TODO for now fixed 0
        BUFFER_TRANSACTION_NEW_FRAG(&trx_buf->hdr_buffer, 0, frag, goto fail_rollback)
        CHECK_BYTES_AVAILABLE(&trx_buf->hdr_buffer, 1, goto fail_rollback)
        *WRITE_POS(frag) = 0;
        DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)

        // Will Topic [MQTT-3.1.3.3]
        CHECK_BYTES_AVAILABLE(&trx_buf->hdr_buffer, 2, goto fail_rollback)
        PACK_2B_INT(&trx_buf->hdr_buffer, will_topic_len, frag)
        if (optimized_add(&trx_buf->hdr_buffer, lwt->will_topic, will_topic_len, lwt->will_topic_free, &frag))
            goto fail_rollback;

        // Will Payload [MQTT-3.1.3.4]
        if (lwt->will_message_size) {
            BUFFER_TRANSACTION_NEW_FRAG(&trx_buf->hdr_buffer, 0, frag, goto fail_rollback)
            CHECK_BYTES_AVAILABLE(&trx_buf->hdr_buffer, 2, goto fail_rollback)
            PACK_2B_INT(&trx_buf->hdr_buffer, lwt->will_message_size, frag)
            if (optimized_add(
                    &trx_buf->hdr_buffer, lwt->will_message, lwt->will_message_size, lwt->will_topic_free, &frag))
                goto fail_rollback;
        }
    }

    // [MQTT-3.1.3.5]
    if (auth->username) {
        BUFFER_TRANSACTION_NEW_FRAG(&trx_buf->hdr_buffer, 0, frag, goto fail_rollback)
        CHECK_BYTES_AVAILABLE(&trx_buf->hdr_buffer, 2, goto fail_rollback)
        PACK_2B_INT(&trx_buf->hdr_buffer, username_len, frag)
        if (optimized_add(&trx_buf->hdr_buffer, auth->username, username_len, auth->username_free, &frag))
            goto fail_rollback;
    }

    // [MQTT-3.1.3.6]
    if (auth->password) {
        BUFFER_TRANSACTION_NEW_FRAG(&trx_buf->hdr_buffer, 0, frag, goto fail_rollback)
        CHECK_BYTES_AVAILABLE(&trx_buf->hdr_buffer, 2, goto fail_rollback)
        PACK_2B_INT(&trx_buf->hdr_buffer, password_len, frag)
        if (optimized_add(&trx_buf->hdr_buffer, auth->password, password_len, auth->password_free, &frag))
            goto fail_rollback;
    }
    trx_buf->hdr_buffer.tail_frag->flags |= BUFFER_FRAG_MQTT_PACKET_TAIL;
    transaction_buffer_transaction_commit(trx_buf)
    return ret;
fail_rollback:
    transaction_buffer_transaction_rollback(trx_buf, ret);
    return NULL;
}

int mqtt_ng_connect(
    struct mqtt_ng_client *client,
    struct mqtt_auth_properties *auth,
    struct mqtt_lwt_properties *lwt,
    uint16_t keep_alive)
{
    client->client_state = MQTT_STATE_RAW;
    mqtt_ng_parser_reset(&client->parser);

    LOCK_HDR_BUFFER(&client->main_buffer);
    client->main_buffer.sending_frag = NULL;
    buffer_purge(&client->main_buffer.hdr_buffer);
    UNLOCK_HDR_BUFFER(&client->main_buffer);

    destroy_timeout_monitor_list(client);

    spinlock_lock(&client->tx_topic_aliases.spinlock);
    // according to MQTT spec topic aliases should not be persisted
    // even if clean session is true
    mqtt_ng_destroy_tx_alias_hash(client->tx_topic_aliases.stoi_dict);

    client->tx_topic_aliases.stoi_dict = TX_ALIASES_INITIALIZE();
    client->tx_topic_aliases.idx_assigned = 0;
    spinlock_unlock(&client->tx_topic_aliases.spinlock);

    mqtt_ng_destroy_rx_alias_hash(client->rx_aliases);
    client->rx_aliases = RX_ALIASES_INITIALIZE();

    client->connect_msg = mqtt_ng_generate_connect(&client->main_buffer, auth, lwt, keep_alive);
    if (client->connect_msg == NULL)
        return 1;

     __atomic_store_n(&client->stats.tx_messages_queued, 1, __ATOMIC_RELAXED);
     __atomic_store_n(&client->stats.tx_messages_sent, 0, __ATOMIC_RELAXED);
     __atomic_store_n(&client->stats.rx_messages_rcvd, 0, __ATOMIC_RELAXED);

    client->client_state = MQTT_STATE_CONNECT_PENDING;
    return 0;
}

uint16_t get_unused_packet_id() {
    static uint16_t packet_id = 0;
    uint16_t id = __atomic_fetch_add(&packet_id, 1, __ATOMIC_RELAXED) + 1;
    return id ? id : 1;
}

static size_t mqtt_ng_publish_size(
    uint16_t topic_len,
    size_t msg_len,
    uint16_t topic_id)
{
    size_t retval = 2
                    + topic_len                           /* Topic Name Length */
                    + 2                                   /* Packet identifier */
                    + 1                                   /* Properties Length for now fixed to 1 property */
                    + msg_len;

    if (topic_id)
        retval += 3;

    return retval;
}

int mqtt_ng_generate_publish(struct transaction_buffer *trx_buf,
                             char *topic,
                             free_fnc_t topic_free,
                             void *msg,
                             free_fnc_t msg_free,
                             size_t msg_len,
                             uint8_t publish_flags,
                             uint16_t *packet_id,
                             uint16_t topic_alias,
                             uint16_t topic_len)
{
    // >> START THE RODEO <<
    transaction_buffer_transaction_start(trx_buf)

    // Calculate the resulting message size sans fixed MQTT header
    size_t size = mqtt_ng_publish_size(topic_len, msg_len, topic_alias);

    // Start generating the message
    struct buffer_fragment *frag = NULL;
    mqtt_msg_data mqtt_msg = NULL;

    BUFFER_TRANSACTION_NEW_FRAG(&trx_buf->hdr_buffer, BUFFER_FRAG_MQTT_PACKET_HEAD, frag, goto fail_rollback )
    // in case of QOS 0 we can garbage collect immediatelly after sending
    uint8_t qos = (publish_flags >> 1) & 0x03;
    if (!qos)
        frag->flags |= BUFFER_FRAG_GARBAGE_COLLECT_ON_SEND;
    mqtt_msg = frag;

    // MQTT Fixed Header
    size_t needed_bytes = 1 /* Packet type */ + MQTT_VARSIZE_INT_BYTES(size) + size - msg_len;
    CHECK_BYTES_AVAILABLE(&trx_buf->hdr_buffer, needed_bytes, goto fail_rollback)

    *WRITE_POS(frag) = (MQTT_CPT_PUBLISH << 4) | (publish_flags & 0xF);
    DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)
    DATA_ADVANCE(&trx_buf->hdr_buffer, uint32_to_mqtt_vbi(size, WRITE_POS(frag)), frag)

    // MQTT Variable Header
    // [MQTT-3.3.2.1]
    PACK_2B_INT(&trx_buf->hdr_buffer, topic_len, frag)
    if (topic != NULL) {
        if (optimized_add(&trx_buf->hdr_buffer, topic, topic_len, topic_free, &frag))
            goto fail_rollback;
        BUFFER_TRANSACTION_NEW_FRAG(&trx_buf->hdr_buffer, 0, frag, goto fail_rollback)
    }

    // [MQTT-3.3.2.2]
    mqtt_msg->packet_id = get_unused_packet_id();
    *packet_id = mqtt_msg->packet_id;
    PACK_2B_INT(&trx_buf->hdr_buffer, mqtt_msg->packet_id, frag)

    // [MQTT-3.3.2.3.1] TODO Property Length for now fixed 0
    *WRITE_POS(frag) = topic_alias ? 3 : 0;
    DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)

    if(topic_alias) {
        *WRITE_POS(frag) = MQTT_PROP_TOPIC_ALIAS;
        DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)

        PACK_2B_INT(&trx_buf->hdr_buffer, topic_alias, frag)
    }

    if( (frag = buffer_new_frag(&trx_buf->hdr_buffer, BUFFER_FRAG_DATA_EXTERNAL)) == NULL )
        goto fail_rollback;

    if (frag_set_external_data(frag, msg, msg_len, msg_free))
        goto fail_rollback;

    trx_buf->hdr_buffer.tail_frag->flags |= BUFFER_FRAG_MQTT_PACKET_TAIL;
    if (!qos)
        trx_buf->hdr_buffer.tail_frag->flags |= BUFFER_FRAG_GARBAGE_COLLECT_ON_SEND;
    transaction_buffer_transaction_commit(trx_buf)
    // mark enqueue time on the HEAD fragment (after commit) so we measure from commit time
    if (mqtt_msg)
        mqtt_msg->enqueued_monotonic_ut = now_monotonic_usec();
    return MQTT_NG_MSGGEN_OK;
fail_rollback:
    transaction_buffer_transaction_rollback(trx_buf, mqtt_msg);
    return MQTT_NG_MSGGEN_BUFFER_OOM;
}

static void mark_message_for_gc(struct buffer_fragment *frag)
{
    while (frag) {
        frag->flags |= BUFFER_FRAG_GARBAGE_COLLECT;
        buffer_frag_free_data(frag);
        if (frag->flags & BUFFER_FRAG_MQTT_PACKET_TAIL)
            return;
        frag = frag->next;
    }
}

// Check if sending_frag points to any fragment in the message starting at msg_head
static bool sending_frag_in_message(struct transaction_buffer *buf, struct buffer_fragment *msg_head)
{
    struct buffer_fragment *frag = msg_head;
    while (frag) {
        if (buf->sending_frag == frag)
            return true;
        if (frag->flags & BUFFER_FRAG_MQTT_PACKET_TAIL)
            return false;
        frag = frag->next;
    }
    return false;
}

static int mark_packet_acked(struct mqtt_ng_client *client, uint16_t packet_id)
{
    size_t reclaimable = 0;
    spinlock_lock(&client->pending_packets.spinlock);
    LOCK_HDR_BUFFER(&client->main_buffer);
    struct buffer_fragment *frag = BUFFER_FIRST_FRAG(&client->main_buffer.hdr_buffer);
    while (frag) {
        if ( (frag->flags & BUFFER_FRAG_MQTT_PACKET_HEAD) && frag->packet_id == packet_id) {
            if (!frag->sent) {
                nd_log(NDLS_DAEMON, NDLP_ERR, "Received packet_id (%" PRIu16 ") belongs to MQTT packet which was not yet sent!", packet_id);
                UNLOCK_HDR_BUFFER(&client->main_buffer);
                spinlock_unlock(&client->pending_packets.spinlock);
                return 1;
            }
            // Do not reprocess this packet
            frag->packet_id = 0;
            usec_t latency = now_monotonic_usec() - frag->sent_monotonic_ut;
            pulse_aclk_sent_message_acked(latency, frag->len);
            __atomic_store_n(&publish_latency, latency, __ATOMIC_RELEASE);

            // Invalidate sending_frag if it points to any fragment in this message
            // since mark_message_for_gc will free the data
            if (sending_frag_in_message(&client->main_buffer, frag))
                client->main_buffer.sending_frag = NULL;

            mark_message_for_gc(frag);

            size_t used = BUFFER_BYTES_USED(&client->main_buffer.hdr_buffer);
            if (reclaimable >= (used / 4))
                transaction_buffer_garbage_collect(&client->main_buffer, true);

            UNLOCK_HDR_BUFFER(&client->main_buffer);
            remove_packet_from_timeout_monitor_list_unsafe(client, packet_id);
            spinlock_unlock(&client->pending_packets.spinlock);
            return 0;
        }

        if(frag_is_marked_for_gc(frag))
            reclaimable += FRAG_SIZE_IN_BUFFER(frag);

        frag = frag->next;
    }
    nd_log(NDLS_DAEMON, NDLP_WARNING, "Received packet_id (%" PRIu16 ") is unknown, removing from monitor list", packet_id);
    UNLOCK_HDR_BUFFER(&client->main_buffer);
    remove_packet_from_timeout_monitor_list_unsafe(client, packet_id);
    spinlock_unlock(&client->pending_packets.spinlock);
    return 1;
}

#define MAX_TIMED_OUT_PACKETS (1024)

static bool check_packet_monitor_list_for_timeouts(struct mqtt_ng_client *client)
{
    uint16_t timed_out_packets[MAX_TIMED_OUT_PACKETS];
    size_t timed_out_count = 0;

    spinlock_lock(&client->pending_packets.spinlock);
    bool first_then_next = true;
    uint32_t *Pvalue;
    Word_t packet_id = 0;
    time_t now = now_realtime_sec();

    while ((Pvalue = (uint32_t *) JudyLFirstThenNext(client->pending_packets.JudyL, &packet_id, &first_then_next))) {
        uint32_t expire_time_delta = *Pvalue;
        if (now >= (PACKET_TIMEOUT_EPOCH + expire_time_delta)) {
            if (timed_out_count < MAX_TIMED_OUT_PACKETS) {
                timed_out_packets[timed_out_count++] = (uint16_t)packet_id;
            } else
                break;
        }
    }
    spinlock_unlock(&client->pending_packets.spinlock);

    // Process timeouts outside the lock
    for (size_t i = 0; i < timed_out_count; i++) {
        mark_packet_acked(client, timed_out_packets[i]);
    }

    return (timed_out_count ==  MAX_TIMED_OUT_PACKETS);
}

#define PUBLISH_SP_SIZE 64
int mqtt_ng_publish(struct mqtt_ng_client *client,
                    char *topic,
                    free_fnc_t topic_free,
                    void *msg,
                    free_fnc_t msg_free,
                    size_t msg_len,
                    uint8_t publish_flags,
                    uint16_t *packet_id)
{
    uint16_t topic_len;
    if (!mqtt_ng_get_2byte_field_length("Topic Name", topic, &topic_len))
        return MQTT_NG_MSGGEN_USER_ERROR;

    struct topic_alias_data *alias = NULL;
    spinlock_lock(&client->tx_topic_aliases.spinlock);
    c_rhash_get_ptr_by_str(client->tx_topic_aliases.stoi_dict, topic, (void**)&alias);
    spinlock_unlock(&client->tx_topic_aliases.spinlock);

    uint16_t topic_id = 0;

    if (alias != NULL) {
        topic_id = alias->idx;
        uint32_t cnt = __atomic_fetch_add(&alias->usage_count, 1, __ATOMIC_SEQ_CST);
        if (cnt) {
            topic = NULL;
            topic_free = NULL;
            topic_len = 0;
        }
    }

    if (client->max_msg_size && PUBLISH_SP_SIZE + mqtt_ng_publish_size(topic_len, msg_len, topic_id) > client->max_msg_size) {
        nd_log(NDLS_DAEMON, NDLP_ERR, "Message too big for server: %zu", msg_len);
        return MQTT_NG_MSGGEN_MSG_TOO_BIG;
    }

    // Ownership contract on failure: do NOT free msg or clear *packet_id here.
    // The sole caller (mqtt_wss_publish5) owns cleanup on every non-OK return.
    // On MQTT_NG_MSGGEN_OK, msg is attached to a buffer fragment and freed by
    // the transaction-buffer GC after ack; *packet_id has been written by the
    // generator. See the long comment in mqtt_wss_publish5 for the invariant.
    int rc = TRY_GENERATE_MESSAGE(
        mqtt_ng_generate_publish,
        topic,
        topic_free,
        msg,
        msg_free,
        msg_len,
        publish_flags,
        packet_id,
        topic_id,
        topic_len);
    if (rc == MQTT_NG_MSGGEN_OK && packet_id)
        add_packet_to_timeout_monitor_list(client, *packet_id);
    return rc;
}

static int mqtt_ng_subscribe_size(struct mqtt_sub *subs, size_t sub_count, size_t *size)
{
    size_t len = 2 /* Packet Identifier */ + 1 /* Properties Length TODO for now fixed 0 */;
    len += sub_count * (2 /* topic filter string length */ + 1 /* [MQTT-3.8.3.1] Subscription Options Byte */);

    for (size_t i = 0; i < sub_count; i++) {
        uint16_t topic_len;
        if (!mqtt_ng_get_2byte_field_length("Topic Filter", subs[i].topic, &topic_len))
            return MQTT_NG_MSGGEN_USER_ERROR;
        len += topic_len;
    }

    *size = len;
    return MQTT_NG_MSGGEN_OK;
}

int mqtt_ng_generate_subscribe(struct transaction_buffer *trx_buf, struct mqtt_sub *subs, size_t sub_count)
{
    size_t size;
    int rc = mqtt_ng_subscribe_size(subs, sub_count, &size);
    if (rc != MQTT_NG_MSGGEN_OK)
        return rc;

    // >> START THE RODEO <<
    transaction_buffer_transaction_start(trx_buf)

    // Start generating the message
    struct buffer_fragment *frag = NULL;
    mqtt_msg_data ret = NULL;

    BUFFER_TRANSACTION_NEW_FRAG(&trx_buf->hdr_buffer, BUFFER_FRAG_MQTT_PACKET_HEAD, frag, goto fail_rollback)
    ret = frag;

    // MQTT Fixed Header
    size_t needed_bytes = 1 /* Packet type */ + MQTT_VARSIZE_INT_BYTES(size) + 3 /*Packet ID + Property Length*/;
    CHECK_BYTES_AVAILABLE(&trx_buf->hdr_buffer, needed_bytes, goto fail_rollback)

    *WRITE_POS(frag) = (MQTT_CPT_SUBSCRIBE << 4) | 0x2 /* [MQTT-3.8.1-1] */;
    DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)
    DATA_ADVANCE(&trx_buf->hdr_buffer, uint32_to_mqtt_vbi(size, WRITE_POS(frag)), frag)

    // MQTT Variable Header
    // [MQTT-3.8.2] PacketID
    ret->packet_id = get_unused_packet_id();
    PACK_2B_INT(&trx_buf->hdr_buffer, ret->packet_id, frag)

    // [MQTT-3.8.2.1.1] Property Length // TODO for now fixed 0
    *WRITE_POS(frag) = 0;
    DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)

    for (size_t i = 0; i < sub_count; i++) {
        uint16_t topic_len;
        if (!mqtt_ng_get_2byte_field_length("Topic Filter", subs[i].topic, &topic_len))
            goto fail_user_rollback;
        BUFFER_TRANSACTION_NEW_FRAG(&trx_buf->hdr_buffer, 0, frag, goto fail_rollback)
        PACK_2B_INT(&trx_buf->hdr_buffer, topic_len, frag)
        if (optimized_add(&trx_buf->hdr_buffer, subs[i].topic, topic_len, subs[i].topic_free, &frag))
            goto fail_rollback;
        BUFFER_TRANSACTION_NEW_FRAG(&trx_buf->hdr_buffer, 0, frag, goto fail_rollback)
        *WRITE_POS(frag) = subs[i].options;
        DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)
    }

    trx_buf->hdr_buffer.tail_frag->flags |= BUFFER_FRAG_MQTT_PACKET_TAIL;
    transaction_buffer_transaction_commit(trx_buf)
    return MQTT_NG_MSGGEN_OK;
fail_user_rollback:
    transaction_buffer_transaction_rollback(trx_buf, ret);
    return MQTT_NG_MSGGEN_USER_ERROR;
fail_rollback:
    transaction_buffer_transaction_rollback(trx_buf, ret);
    return MQTT_NG_MSGGEN_BUFFER_OOM;
}

int mqtt_ng_subscribe(struct mqtt_ng_client *client, struct mqtt_sub *subs, size_t sub_count)
{
    return TRY_GENERATE_MESSAGE(mqtt_ng_generate_subscribe, subs, sub_count);
}

int mqtt_ng_generate_disconnect(struct transaction_buffer *trx_buf, uint8_t reason_code)
{
    // >> START THE RODEO <<
    transaction_buffer_transaction_start(trx_buf)

    // Calculate the resulting message size sans fixed MQTT header
    size_t size = reason_code ? 1 : 0;

    // Start generating the message
    struct buffer_fragment *frag = NULL;
    mqtt_msg_data ret = NULL;

    BUFFER_TRANSACTION_NEW_FRAG(&trx_buf->hdr_buffer, BUFFER_FRAG_MQTT_PACKET_HEAD, frag, goto fail_rollback)
    ret = frag;

    // MQTT Fixed Header
    size_t needed_bytes = 1 /* Packet type */ + MQTT_VARSIZE_INT_BYTES(size) + (reason_code ? 1 : 0);
    CHECK_BYTES_AVAILABLE(&trx_buf->hdr_buffer, needed_bytes, goto fail_rollback)

    *WRITE_POS(frag) = MQTT_CPT_DISCONNECT << 4;
    DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)
    DATA_ADVANCE(&trx_buf->hdr_buffer, uint32_to_mqtt_vbi(size, WRITE_POS(frag)), frag)

    if (reason_code) {
        // MQTT Variable Header
        // [MQTT-3.14.2.1] PacketID
        *WRITE_POS(frag) = reason_code;
        DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)
    }

    trx_buf->hdr_buffer.tail_frag->flags |= BUFFER_FRAG_MQTT_PACKET_TAIL;
    transaction_buffer_transaction_commit(trx_buf)
    return MQTT_NG_MSGGEN_OK;
fail_rollback:
    transaction_buffer_transaction_rollback(trx_buf, ret);
    return MQTT_NG_MSGGEN_BUFFER_OOM;
}

int mqtt_ng_disconnect(struct mqtt_ng_client *client, uint8_t reason_code)
{
    return TRY_GENERATE_MESSAGE(mqtt_ng_generate_disconnect, reason_code);
}

static int mqtt_generate_puback(struct transaction_buffer *trx_buf, uint16_t packet_id, uint8_t reason_code)
{
    // >> START THE RODEO <<
    transaction_buffer_transaction_start(trx_buf)

    // Calculate the resulting message size sans fixed MQTT header
    size_t size = 2 /* Packet ID */ + (reason_code ? 1 : 0) /* reason code */;

    // Start generating the message
    struct buffer_fragment *frag = NULL;

    BUFFER_TRANSACTION_NEW_FRAG(&trx_buf->hdr_buffer, BUFFER_FRAG_MQTT_PACKET_HEAD | BUFFER_FRAG_GARBAGE_COLLECT_ON_SEND, frag, goto fail_rollback)

    // MQTT Fixed Header
    size_t needed_bytes = 1 /* Packet type */ + MQTT_VARSIZE_INT_BYTES(size) + size;
    CHECK_BYTES_AVAILABLE(&trx_buf->hdr_buffer, needed_bytes, goto fail_rollback)

    *WRITE_POS(frag) = MQTT_CPT_PUBACK << 4;
    DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)
    DATA_ADVANCE(&trx_buf->hdr_buffer, uint32_to_mqtt_vbi(size, WRITE_POS(frag)), frag)

    // MQTT Variable Header
    PACK_2B_INT(&trx_buf->hdr_buffer, packet_id, frag)

    if (reason_code) {
        // MQTT Variable Header
        // [MQTT-3.14.2.1] PacketID
        *WRITE_POS(frag) = reason_code;
        DATA_ADVANCE(&trx_buf->hdr_buffer, 1, frag)
    }

    trx_buf->hdr_buffer.tail_frag->flags |= BUFFER_FRAG_MQTT_PACKET_TAIL;
    transaction_buffer_transaction_commit(trx_buf)
    return MQTT_NG_MSGGEN_OK;
fail_rollback:
    transaction_buffer_transaction_rollback(trx_buf, frag);
    return MQTT_NG_MSGGEN_BUFFER_OOM;
}

static int mqtt_ng_puback(struct mqtt_ng_client *client, uint16_t packet_id, uint8_t reason_code)
{
    return TRY_GENERATE_MESSAGE(mqtt_generate_puback, packet_id, reason_code);
}

int mqtt_ng_ping(struct mqtt_ng_client *client)
{
    client->ping_pending = 1;
    return MQTT_NG_MSGGEN_OK;
}

#define MQTT_NG_CLIENT_NEED_MORE_BYTES         0x10
#define MQTT_NG_CLIENT_MQTT_PACKET_DONE        0x11
#define MQTT_NG_CLIENT_PARSE_DONE              0x12
#define MQTT_NG_CLIENT_WANT_WRITE              0x13
#define MQTT_NG_CLIENT_OK_CALL_AGAIN           0
#define MQTT_NG_CLIENT_PROTOCOL_ERROR         (-1)
#define MQTT_NG_CLIENT_SERVER_RETURNED_ERROR  (-2)
#define MQTT_NG_CLIENT_NOT_IMPL_YET           (-3)
#define MQTT_NG_CLIENT_INTERNAL_ERROR         (-5)

#define BUF_READ_CHECK_AT_LEAST(buf, x)                                                                                \
    if (rbuf_bytes_available(buf) < (x))                                                                               \
        return MQTT_NG_CLIENT_NEED_MORE_BYTES;

static int mqtt_properties_read_check(struct mqtt_properties_parser_ctx *ctx, rbuf_t data, size_t bytes)
{
    size_t consumed = ctx->bytes_consumed;
    if (ctx->state == PROPERTIES_LENGTH || ctx->state == PROPERTY_TYPE_VBI)
        consumed += ctx->vbi_parser_ctx.bytes;

    if (consumed > ctx->max_bytes || bytes > ctx->max_bytes - consumed) {
        nd_log(NDLS_DAEMON, NDLP_ERR, "MQTT properties exceed packet boundary.");
        return MQTT_NG_CLIENT_PROTOCOL_ERROR;
    }
    if (rbuf_bytes_available(data) < bytes)
        return MQTT_NG_CLIENT_NEED_MORE_BYTES;

    return MQTT_NG_CLIENT_PARSE_DONE;
}

#define PROPERTIES_READ_CHECK_AT_LEAST(ctx, data, x) do {                                                              \
        int read_check_rc = mqtt_properties_read_check((ctx), (data), (x));                                            \
        if (read_check_rc != MQTT_NG_CLIENT_PARSE_DONE)                                                               \
            return read_check_rc;                                                                                     \
    } while(0)

#define vbi_parser_reset_ctx(ctx) memset(ctx, 0, sizeof(struct mqtt_vbi_parser_ctx))

static int vbi_parser_parse(struct mqtt_vbi_parser_ctx *ctx, rbuf_t data)
{
    if (ctx->bytes > MQTT_VBI_MAXBYTES - 1) {
        nd_log(NDLS_DAEMON, NDLP_ERR, "MQTT Variable Byte Integer can't be longer than %d bytes", MQTT_VBI_MAXBYTES);
        return MQTT_NG_CLIENT_PROTOCOL_ERROR;
    }
    if (!ctx->bytes || ctx->data[ctx->bytes-1] & MQTT_VBI_CONTINUATION_FLAG) {
        BUF_READ_CHECK_AT_LEAST(data, 1)
        ctx->bytes++;
        rbuf_pop(data, &ctx->data[ctx->bytes-1], 1);
        if ( ctx->data[ctx->bytes-1] & MQTT_VBI_CONTINUATION_FLAG )
            return MQTT_NG_CLIENT_OK_CALL_AGAIN;
    }

    if (mqtt_vbi_to_uint32(ctx->data, &ctx->result)) {
            nd_log(NDLS_DAEMON, NDLP_ERR, "MQTT Variable Byte Integer failed to be parsed.");
            return MQTT_NG_CLIENT_PROTOCOL_ERROR;
    }

    return MQTT_NG_CLIENT_PARSE_DONE;
}

static void mqtt_properties_parser_ctx_reset(struct mqtt_properties_parser_ctx *ctx)
{
    ctx->state = PROPERTIES_LENGTH;
    while (ctx->head) {
        struct mqtt_property *f = ctx->head;
        ctx->head = ctx->head->next;
        if (f->type == MQTT_TYPE_STR || f->type == MQTT_TYPE_STR_PAIR)
            freez(f->data.strings[0]);
        if (f->type == MQTT_TYPE_STR_PAIR)
            freez(f->data.strings[1]);
        if (f->type == MQTT_TYPE_BIN)
            freez(f->data.bindata);
        freez(f);
    }
    ctx->tail = NULL;
    ctx->properties_length = 0;
    ctx->bytes_consumed = 0;
    ctx->max_bytes = SIZE_MAX;
    vbi_parser_reset_ctx(&ctx->vbi_parser_ctx);
}

static void mqtt_ng_parser_reset(struct mqtt_ng_parser *parser)
{
    rbuf_t received_data = parser->received_data;
    uint8_t control_packet_type = get_control_packet_type(parser->mqtt_control_packet_type);
    bool current_variable_header =
        parser->state == MQTT_PARSE_VARIABLE_HEADER &&
        parser->varhdr_state != MQTT_PARSE_VARHDR_INITIAL;

    mqtt_properties_parser_ctx_reset(&parser->properties_parser);

    // VARHDR_INITIAL may still contain stale pointers from an already handled packet.
    // Incomplete PUBLISH state can own the topic, but payload data is allocated
    // only after the full payload is available.
    if (current_variable_header && control_packet_type == MQTT_CPT_PUBLISH)
        freez(parser->mqtt_packet.publish.topic);

    if (current_variable_header && control_packet_type == MQTT_CPT_SUBACK)
        freez(parser->mqtt_packet.suback.reason_codes);

    memset(parser, 0, sizeof(*parser));
    parser->received_data = received_data;
    parser->state = MQTT_PARSE_FIXED_HEADER_PACKET_TYPE;
}

struct mqtt_property_type {
    uint8_t id;
    enum mqtt_datatype datatype;
    const char* name;
};

const struct mqtt_property_type mqtt_property_types[] = {
    { .id = MQTT_PROP_TOPIC_ALIAS,             .name = MQTT_PROP_TOPIC_ALIAS_NAME,             .datatype = MQTT_TYPE_UINT_16  },

    { .id = MQTT_PROP_PAYLOAD_FMT_INDICATOR,   .name = MQTT_PROP_PAYLOAD_FMT_INDICATOR_NAME,   .datatype = MQTT_TYPE_UINT_8   },
    { .id = MQTT_PROP_MSG_EXPIRY_INTERVAL,     .name = MQTT_PROP_MSG_EXPIRY_INTERVAL_NAME,     .datatype = MQTT_TYPE_UINT_32  },
    { .id = MQTT_PROP_CONTENT_TYPE,            .name = MQTT_PROP_CONTENT_TYPE_NAME,            .datatype = MQTT_TYPE_STR      },
    { .id = MQTT_PROP_RESPONSE_TOPIC,          .name = MQTT_PROP_RESPONSE_TOPIC_NAME,          .datatype = MQTT_TYPE_STR      },
    { .id = MQTT_PROP_CORRELATION_DATA,        .name = MQTT_PROP_CORRELATION_DATA_NAME,        .datatype = MQTT_TYPE_BIN      },
    { .id = MQTT_PROP_SUB_IDENTIFIER,          .name = MQTT_PROP_SUB_IDENTIFIER_NAME,          .datatype = MQTT_TYPE_VBI      },
    { .id = MQTT_PROP_SESSION_EXPIRY_INTERVAL, .name = MQTT_PROP_SESSION_EXPIRY_INTERVAL_NAME, .datatype = MQTT_TYPE_UINT_32  },
    { .id = MQTT_PROP_ASSIGNED_CLIENT_ID,      .name = MQTT_PROP_ASSIGNED_CLIENT_ID_NAME,      .datatype = MQTT_TYPE_STR      },
    { .id = MQTT_PROP_SERVER_KEEP_ALIVE,       .name = MQTT_PROP_SERVER_KEEP_ALIVE_NAME,       .datatype = MQTT_TYPE_UINT_16  },
    { .id = MQTT_PROP_AUTH_METHOD,             .name = MQTT_PROP_AUTH_METHOD_NAME,             .datatype = MQTT_TYPE_STR      },
    { .id = MQTT_PROP_AUTH_DATA,               .name = MQTT_PROP_AUTH_DATA_NAME,               .datatype = MQTT_TYPE_BIN      },
    { .id = MQTT_PROP_REQ_PROBLEM_INFO,        .name = MQTT_PROP_REQ_PROBLEM_INFO_NAME,        .datatype = MQTT_TYPE_UINT_8   },
    { .id = MQTT_PROP_WILL_DELAY_INTERVAL,     .name = MQTT_PROP_WIIL_DELAY_INTERVAL_NAME,     .datatype = MQTT_TYPE_UINT_32  },
    { .id = MQTT_PROP_REQ_RESP_INFORMATION,    .name = MQTT_PROP_REQ_RESP_INFORMATION_NAME,    .datatype = MQTT_TYPE_UINT_8   },
    { .id = MQTT_PROP_RESP_INFORMATION,        .name = MQTT_PROP_RESP_INFORMATION_NAME,        .datatype = MQTT_TYPE_STR      },
    { .id = MQTT_PROP_SERVER_REF,              .name = MQTT_PROP_SERVER_REF_NAME,              .datatype = MQTT_TYPE_STR      },
    { .id = MQTT_PROP_REASON_STR,              .name = MQTT_PROP_REASON_STR_NAME,              .datatype = MQTT_TYPE_STR      },
    { .id = MQTT_PROP_RECEIVE_MAX,             .name = MQTT_PROP_RECEIVE_MAX_NAME,             .datatype = MQTT_TYPE_UINT_16  },
    { .id = MQTT_PROP_TOPIC_ALIAS_MAX,         .name = MQTT_PROP_TOPIC_ALIAS_MAX_NAME,         .datatype = MQTT_TYPE_UINT_16  },
    // MQTT_PROP_TOPIC_ALIAS is first as it is most often used
    { .id = MQTT_PROP_MAX_QOS,                 .name = MQTT_PROP_MAX_QOS_NAME,                 .datatype = MQTT_TYPE_UINT_8   },
    { .id = MQTT_PROP_RETAIN_AVAIL,            .name = MQTT_PROP_RETAIN_AVAIL_NAME,            .datatype = MQTT_TYPE_UINT_8   },
    { .id = MQTT_PROP_USR,                     .name = MQTT_PROP_USR_NAME,                     .datatype = MQTT_TYPE_STR_PAIR },
    { .id = MQTT_PROP_MAX_PKT_SIZE,            .name = MQTT_PROP_MAX_PKT_SIZE_NAME,            .datatype = MQTT_TYPE_UINT_32  },
    { .id = MQTT_PROP_WILDCARD_SUB_AVAIL,      .name = MQTT_PROP_WILDCARD_SUB_AVAIL_NAME,      .datatype = MQTT_TYPE_UINT_8   },
    { .id = MQTT_PROP_SUB_ID_AVAIL,            .name = MQTT_PROP_SUB_ID_AVAIL_NAME,            .datatype = MQTT_TYPE_UINT_8   },
    { .id = MQTT_PROP_SHARED_SUB_AVAIL,        .name = MQTT_PROP_SHARED_SUB_AVAIL_NAME,        .datatype = MQTT_TYPE_UINT_8   },
    { .id = 0,                                 .name = NULL,                                   .datatype = MQTT_TYPE_UNKNOWN  }
};

static int get_property_type_by_id(uint8_t property_id) {
    for (int i = 0; mqtt_property_types[i].datatype != MQTT_TYPE_UNKNOWN; i++) {
        if (mqtt_property_types[i].id == property_id)
            return mqtt_property_types[i].datatype;
    }
    return MQTT_TYPE_UNKNOWN;
}

struct mqtt_property *get_property_by_id(struct mqtt_property *props, uint8_t property_id)
{
    while (props) {
        if (props->id == property_id) {
            return props;
        }
        props = props->next;
    }
    return NULL;
}

// Parses [MQTT-2.2.2]
static int parse_properties_array(struct mqtt_properties_parser_ctx *ctx, rbuf_t data)
{
    int rc;
    switch (ctx->state) {
        case PROPERTIES_LENGTH:
            PROPERTIES_READ_CHECK_AT_LEAST(ctx, data, 1);
            rc = vbi_parser_parse(&ctx->vbi_parser_ctx, data);
            if (rc == MQTT_NG_CLIENT_PARSE_DONE) {
                ctx->properties_length = ctx->vbi_parser_ctx.result;
                ctx->bytes_consumed += ctx->vbi_parser_ctx.bytes;
                ctx->vbi_length = ctx->vbi_parser_ctx.bytes;
                if (ctx->vbi_length > ctx->max_bytes || ctx->properties_length > ctx->max_bytes - ctx->vbi_length) {
                    nd_log(NDLS_DAEMON, NDLP_ERR, "MQTT properties exceed packet boundary.");
                    return MQTT_NG_CLIENT_PROTOCOL_ERROR;
                }
                if (!ctx->properties_length)
                    return MQTT_NG_CLIENT_PARSE_DONE;
                ctx->state = PROPERTY_CREATE;
                break;
            }
            return rc;
        case PROPERTY_CREATE:
            PROPERTIES_READ_CHECK_AT_LEAST(ctx, data, 1);
            struct mqtt_property *prop = callocz(1, sizeof(struct mqtt_property));
            if (ctx->head == NULL) {
                ctx->head = prop;
                ctx->tail = prop;
            } else {
                ctx->tail->next = prop;
                ctx->tail = ctx->tail->next;
            }
            ctx->state = PROPERTY_ID;
            /* FALLTHROUGH */
        case PROPERTY_ID:
            rbuf_pop(data, (char*)&ctx->tail->id, 1);
            ctx->bytes_consumed += 1;
            ctx->tail->type = get_property_type_by_id(ctx->tail->id);
            switch (ctx->tail->type) {
                case MQTT_TYPE_UINT_16:
                    ctx->state = PROPERTY_TYPE_UINT16;
                    break;
                case MQTT_TYPE_UINT_32:
                    ctx->state = PROPERTY_TYPE_UINT32;
                    break;
                case MQTT_TYPE_UINT_8:
                    ctx->state = PROPERTY_TYPE_UINT8;
                    break;
                case MQTT_TYPE_VBI:
                    ctx->state = PROPERTY_TYPE_VBI;
                    vbi_parser_reset_ctx(&ctx->vbi_parser_ctx);
                    break;
                case MQTT_TYPE_STR:
                case MQTT_TYPE_STR_PAIR:
                    ctx->str_idx = 0;
                    /* FALLTHROUGH */
                case MQTT_TYPE_BIN:
                    ctx->state = PROPERTY_TYPE_STR_BIN_LEN;
                    break;
                default:
                    nd_log(NDLS_DAEMON, NDLP_ERR, "Unsupported property type %d for property id %d.", (int)ctx->tail->type, (int)ctx->tail->id);
                    return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            }
            break;
        case PROPERTY_TYPE_STR_BIN_LEN:
            PROPERTIES_READ_CHECK_AT_LEAST(ctx, data, sizeof(uint16_t));
            rbuf_pop(data, (char*)&ctx->tail->bindata_len, sizeof(uint16_t));
            ctx->tail->bindata_len = be16toh(ctx->tail->bindata_len);
            ctx->bytes_consumed += 2;
            switch (ctx->tail->type) {
                case MQTT_TYPE_BIN:
                    ctx->state = PROPERTY_TYPE_BIN;
                    break;
                case MQTT_TYPE_STR:
                case MQTT_TYPE_STR_PAIR:
                    ctx->state = PROPERTY_TYPE_STR;
                    break;
                default:
                    nd_log(NDLS_DAEMON, NDLP_ERR, "Unexpected datatype in PROPERTY_TYPE_STR_BIN_LEN %d", (int)ctx->tail->type);
                    return MQTT_NG_CLIENT_INTERNAL_ERROR;
            }
            break;
        case PROPERTY_TYPE_STR:
            PROPERTIES_READ_CHECK_AT_LEAST(ctx, data, ctx->tail->bindata_len);
            ctx->tail->data.strings[ctx->str_idx] = mallocz(ctx->tail->bindata_len + 1);
            rbuf_pop(data, ctx->tail->data.strings[ctx->str_idx], ctx->tail->bindata_len);
            ctx->tail->data.strings[ctx->str_idx][ctx->tail->bindata_len] = 0;
            ctx->str_idx++;
            ctx->bytes_consumed += ctx->tail->bindata_len;
            if (ctx->tail->type == MQTT_TYPE_STR_PAIR && ctx->str_idx < 2) {
                ctx->state = PROPERTY_TYPE_STR_BIN_LEN;
                break;
            }
            ctx->state = PROPERTY_NEXT;
            break;
        case PROPERTY_TYPE_BIN:
            PROPERTIES_READ_CHECK_AT_LEAST(ctx, data, ctx->tail->bindata_len);
            ctx->tail->data.bindata = mallocz(ctx->tail->bindata_len);
            rbuf_pop(data, ctx->tail->data.bindata, ctx->tail->bindata_len);
            ctx->bytes_consumed += ctx->tail->bindata_len;
            ctx->state = PROPERTY_NEXT;
            break;
        case PROPERTY_TYPE_VBI:
            PROPERTIES_READ_CHECK_AT_LEAST(ctx, data, 1);
            rc = vbi_parser_parse(&ctx->vbi_parser_ctx, data);
            if (rc == MQTT_NG_CLIENT_PARSE_DONE) {
                ctx->tail->data.uint32 = ctx->vbi_parser_ctx.result;
                ctx->bytes_consumed += ctx->vbi_parser_ctx.bytes;
                ctx->state = PROPERTY_NEXT;
                break;
            }
            return rc;
        case PROPERTY_TYPE_UINT8:
            PROPERTIES_READ_CHECK_AT_LEAST(ctx, data, sizeof(uint8_t));
            rbuf_pop(data, (char*)&ctx->tail->data.uint8, sizeof(uint8_t));
            ctx->bytes_consumed += sizeof(uint8_t);
            ctx->state = PROPERTY_NEXT;
            break;
        case PROPERTY_TYPE_UINT32:
            PROPERTIES_READ_CHECK_AT_LEAST(ctx, data, sizeof(uint32_t));
            rbuf_pop(data, (char*)&ctx->tail->data.uint32, sizeof(uint32_t));
            ctx->tail->data.uint32 = be32toh(ctx->tail->data.uint32);
            ctx->bytes_consumed += sizeof(uint32_t);
            ctx->state = PROPERTY_NEXT;
            break;
        case PROPERTY_TYPE_UINT16:
            PROPERTIES_READ_CHECK_AT_LEAST(ctx, data, sizeof(uint16_t));
            rbuf_pop(data, (char*)&ctx->tail->data.uint16, sizeof(uint16_t));
            ctx->tail->data.uint16 = be16toh(ctx->tail->data.uint16);
            ctx->bytes_consumed += sizeof(uint16_t);
            ctx->state = PROPERTY_NEXT;
            /* FALLTHROUGH */
        case PROPERTY_NEXT:
            if (ctx->properties_length > ctx->bytes_consumed - ctx->vbi_length) {
                ctx->state = PROPERTY_CREATE;
                break;
            }
            if (ctx->properties_length < ctx->bytes_consumed - ctx->vbi_length) {
                nd_log(NDLS_DAEMON, NDLP_ERR, "MQTT property exceeds properties length.");
                return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            }
            return MQTT_NG_CLIENT_PARSE_DONE;
    }
    return MQTT_NG_CLIENT_OK_CALL_AGAIN;
}

static int parse_connack_varhdr(struct mqtt_ng_client *client)
{
    int rc;
    struct mqtt_ng_parser *parser = &client->parser;
    switch (parser->varhdr_state) {
        case MQTT_PARSE_VARHDR_INITIAL:
            if (parser->mqtt_fixed_hdr_remaining_length < 3) {
                nd_log(NDLS_DAEMON, NDLP_ERR, "Error parsing CONNACK message");
                return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            }
            BUF_READ_CHECK_AT_LEAST(parser->received_data, 2)
            rbuf_pop(parser->received_data, (char*)&parser->mqtt_packet.connack.flags, 1);
            rbuf_pop(parser->received_data, (char*)&parser->mqtt_packet.connack.reason_code, 1);
            parser->varhdr_state = MQTT_PARSE_VARHDR_PROPS;
            mqtt_properties_parser_ctx_reset(&parser->properties_parser);
            parser->properties_parser.max_bytes = parser->mqtt_fixed_hdr_remaining_length - 2;
            break;
        case MQTT_PARSE_VARHDR_PROPS:
            rc = parse_properties_array(&parser->properties_parser, parser->received_data);
            if (rc != MQTT_NG_CLIENT_PARSE_DONE)
                return rc;
            if (parser->properties_parser.bytes_consumed != parser->properties_parser.max_bytes) {
                nd_log(NDLS_DAEMON, NDLP_ERR, "Error parsing CONNACK message");
                return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            }
            return MQTT_NG_CLIENT_PARSE_DONE;
        default:
            nd_log(NDLS_DAEMON, NDLP_ERR, "invalid state for connack varhdr parser");
            return MQTT_NG_CLIENT_INTERNAL_ERROR;
    }
    return MQTT_NG_CLIENT_OK_CALL_AGAIN;
}

static int parse_disconnect_varhdr(struct mqtt_ng_client *client)
{
    int rc;
    struct mqtt_ng_parser *parser = &client->parser;
    switch (parser->varhdr_state) {
        case MQTT_PARSE_VARHDR_INITIAL:
            if (!parser->mqtt_fixed_hdr_remaining_length) {
                // [MQTT-3.14.2.1] if reason code omitted act same as == 0
                parser->mqtt_packet.disconnect.reason_code = 0;
                return MQTT_NG_CLIENT_PARSE_DONE;
            }
            BUF_READ_CHECK_AT_LEAST(parser->received_data, 1)
            rbuf_pop(parser->received_data, (char*)&parser->mqtt_packet.disconnect.reason_code, 1);
            if (parser->mqtt_fixed_hdr_remaining_length == 1)
                return MQTT_NG_CLIENT_PARSE_DONE;
            parser->varhdr_state = MQTT_PARSE_VARHDR_PROPS;
            mqtt_properties_parser_ctx_reset(&parser->properties_parser);
            parser->properties_parser.max_bytes = parser->mqtt_fixed_hdr_remaining_length - 1;
            break;
        case MQTT_PARSE_VARHDR_PROPS:
            rc = parse_properties_array(&parser->properties_parser, parser->received_data);
            if (rc != MQTT_NG_CLIENT_PARSE_DONE)
                return rc;
            if (parser->properties_parser.bytes_consumed != parser->properties_parser.max_bytes) {
                nd_log(NDLS_DAEMON, NDLP_ERR, "Error parsing DISCONNECT message");
                return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            }
            return MQTT_NG_CLIENT_PARSE_DONE;
        default:
            nd_log(NDLS_DAEMON, NDLP_ERR, "invalid state for connack varhdr parser");
            return MQTT_NG_CLIENT_INTERNAL_ERROR;
    }
    return MQTT_NG_CLIENT_OK_CALL_AGAIN;
}

static int parse_puback_varhdr(struct mqtt_ng_client *client)
{
    int rc;
    struct mqtt_ng_parser *parser = &client->parser;
    switch (parser->varhdr_state) {
        case MQTT_PARSE_VARHDR_INITIAL:
            if (parser->mqtt_fixed_hdr_remaining_length < 2) {
                nd_log(NDLS_DAEMON, NDLP_ERR, "Error parsing PUBACK message");
                return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            }
            BUF_READ_CHECK_AT_LEAST(parser->received_data, 2)
            rbuf_pop(parser->received_data, (char*)&parser->mqtt_packet.puback.packet_id, 2);
            parser->mqtt_packet.puback.packet_id = be16toh(parser->mqtt_packet.puback.packet_id);
            if (parser->mqtt_fixed_hdr_remaining_length < 3) {
                // [MQTT-3.4.2.1] if length is not big enough for reason code
                // it is omitted and handled same as if it was present and == 0
                // initially missed this detail and was wondering WTF is going on (sigh)
                parser->mqtt_packet.puback.reason_code = 0;
                return MQTT_NG_CLIENT_PARSE_DONE;
            }
            parser->varhdr_state = MQTT_PARSE_VARHDR_OPTIONAL_REASON_CODE;
            /* FALLTHROUGH */
        case MQTT_PARSE_VARHDR_OPTIONAL_REASON_CODE:
            BUF_READ_CHECK_AT_LEAST(parser->received_data, 1)
            rbuf_pop(parser->received_data, (char*)&parser->mqtt_packet.puback.reason_code, 1);
            // LOL so in CONNACK you have to have 0 byte to
            // signify empty properties list
            // but in PUBACK it can be omitted if remaining length doesn't allow it (sigh)
            if (parser->mqtt_fixed_hdr_remaining_length < 4)
                return MQTT_NG_CLIENT_PARSE_DONE;

            parser->varhdr_state = MQTT_PARSE_VARHDR_PROPS;
            mqtt_properties_parser_ctx_reset(&parser->properties_parser);
            parser->properties_parser.max_bytes = parser->mqtt_fixed_hdr_remaining_length - 3;
            /* FALLTHROUGH */
        case MQTT_PARSE_VARHDR_PROPS:
            rc = parse_properties_array(&parser->properties_parser, parser->received_data);
            if (rc != MQTT_NG_CLIENT_PARSE_DONE)
                return rc;
            if (parser->properties_parser.bytes_consumed != parser->properties_parser.max_bytes) {
                nd_log(NDLS_DAEMON, NDLP_ERR, "Error parsing PUBACK message");
                return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            }
            return MQTT_NG_CLIENT_PARSE_DONE;
        default:
            nd_log(NDLS_DAEMON, NDLP_ERR, "invalid state for puback varhdr parser");
            return MQTT_NG_CLIENT_INTERNAL_ERROR;
    }
    return MQTT_NG_CLIENT_OK_CALL_AGAIN;
}

static int parse_suback_varhdr(struct mqtt_ng_client *client)
{
    int rc;
    size_t avail, stored;
    struct mqtt_ng_parser *parser = &client->parser;
    struct mqtt_suback *suback = &client->parser.mqtt_packet.suback;
    switch (parser->varhdr_state) {
        case MQTT_PARSE_VARHDR_INITIAL:
            suback->reason_codes = NULL;
            if (parser->mqtt_fixed_hdr_remaining_length < 4) {
                nd_log(NDLS_DAEMON, NDLP_ERR, "Error parsing SUBACK message");
                return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            }
            BUF_READ_CHECK_AT_LEAST(parser->received_data, 2)
            rbuf_pop(parser->received_data, (char*)&suback->packet_id, 2);
            suback->packet_id = be16toh(suback->packet_id);
            parser->varhdr_state = MQTT_PARSE_VARHDR_PROPS;
            parser->mqtt_parsed_len = 2;
            mqtt_properties_parser_ctx_reset(&parser->properties_parser);
            parser->properties_parser.max_bytes = parser->mqtt_fixed_hdr_remaining_length - parser->mqtt_parsed_len - 1;
            /* FALLTHROUGH */
        case MQTT_PARSE_VARHDR_PROPS:
            rc = parse_properties_array(&parser->properties_parser, parser->received_data);
            if (rc != MQTT_NG_CLIENT_PARSE_DONE) 
                return rc;
            parser->mqtt_parsed_len += parser->properties_parser.bytes_consumed;
            suback->reason_code_count = parser->mqtt_fixed_hdr_remaining_length - parser->mqtt_parsed_len;
            if (!suback->reason_code_count) {
                nd_log(NDLS_DAEMON, NDLP_ERR, "Error parsing SUBACK message");
                return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            }
            suback->reason_codes = callocz(suback->reason_code_count, sizeof(*suback->reason_codes));
            suback->reason_codes_pending = suback->reason_code_count;
            parser->varhdr_state = MQTT_PARSE_REASONCODES;
            /* FALLTHROUGH */
        case MQTT_PARSE_REASONCODES:
            avail = rbuf_bytes_available(parser->received_data);
            if (avail < 1)
                return MQTT_NG_CLIENT_NEED_MORE_BYTES;

            stored = suback->reason_code_count - suback->reason_codes_pending;
            suback->reason_codes_pending -= rbuf_pop(
                parser->received_data,
                (char *)&suback->reason_codes[stored],
                MIN(suback->reason_codes_pending, avail));

            if (!suback->reason_codes_pending)
                return MQTT_NG_CLIENT_PARSE_DONE;

            return MQTT_NG_CLIENT_NEED_MORE_BYTES;
        default:
            nd_log(NDLS_DAEMON, NDLP_ERR, "invalid state for suback varhdr parser");
            return MQTT_NG_CLIENT_INTERNAL_ERROR;
    }
    return MQTT_NG_CLIENT_OK_CALL_AGAIN;
}

static int parse_publish_varhdr(struct mqtt_ng_client *client)
{
    int rc;
    struct mqtt_ng_parser *parser = &client->parser;
    struct mqtt_publish *publish = &client->parser.mqtt_packet.publish;
    switch (parser->varhdr_state) {
        case MQTT_PARSE_VARHDR_INITIAL:
            BUF_READ_CHECK_AT_LEAST(parser->received_data, 2)
            publish->topic = NULL;
            publish->data = NULL;
            publish->qos = ((parser->mqtt_control_packet_type >> 1) & 0x03);
            rbuf_pop(parser->received_data, (char*)&publish->topic_len, 2);
            publish->topic_len = be16toh(publish->topic_len);
            parser->mqtt_parsed_len = 2;
            if (parser->mqtt_fixed_hdr_remaining_length < parser->mqtt_parsed_len + publish->topic_len + (publish->qos ? 2 : 0) + 1) {
                nd_log(NDLS_DAEMON, NDLP_ERR, "Error parsing PUBLISH message");
                return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            }
            if (!publish->topic_len) {
                parser->varhdr_state = MQTT_PARSE_VARHDR_POST_TOPICNAME;
                break;
            }
            publish->topic = callocz(1, publish->topic_len + 1 /* add 0x00 */);
            parser->varhdr_state = MQTT_PARSE_VARHDR_TOPICNAME;
            /* FALLTHROUGH */
        case MQTT_PARSE_VARHDR_TOPICNAME:
            // TODO check empty topic can be valid? In which case we have to skip this step
            BUF_READ_CHECK_AT_LEAST(parser->received_data, publish->topic_len)
            rbuf_pop(parser->received_data, publish->topic, publish->topic_len);
            parser->mqtt_parsed_len += publish->topic_len;
            parser->varhdr_state = MQTT_PARSE_VARHDR_POST_TOPICNAME;
            /* FALLTHROUGH */
        case MQTT_PARSE_VARHDR_POST_TOPICNAME:
            mqtt_properties_parser_ctx_reset(&parser->properties_parser);
            if (!publish->qos) { // PacketID present only for QOS > 0 [MQTT-3.3.2.2]
                parser->varhdr_state = MQTT_PARSE_VARHDR_PROPS;
                break;
            }
            parser->varhdr_state = MQTT_PARSE_VARHDR_PACKET_ID;
            /* FALLTHROUGH */
        case MQTT_PARSE_VARHDR_PACKET_ID:
            BUF_READ_CHECK_AT_LEAST(parser->received_data, 2)
            rbuf_pop(parser->received_data, (char*)&publish->packet_id, 2);
            publish->packet_id = be16toh(publish->packet_id);
            parser->varhdr_state = MQTT_PARSE_VARHDR_PROPS;
            parser->mqtt_parsed_len += 2;
            /* FALLTHROUGH */
        case MQTT_PARSE_VARHDR_PROPS:
            parser->properties_parser.max_bytes = parser->mqtt_fixed_hdr_remaining_length - parser->mqtt_parsed_len;
            rc = parse_properties_array(&parser->properties_parser, parser->received_data);
            if (rc != MQTT_NG_CLIENT_PARSE_DONE) 
                return rc;
            parser->mqtt_parsed_len += parser->properties_parser.bytes_consumed;
            parser->varhdr_state = MQTT_PARSE_PAYLOAD;
            /* FALLTHROUGH */
        case MQTT_PARSE_PAYLOAD:
            if (parser->mqtt_fixed_hdr_remaining_length < parser->mqtt_parsed_len) {
                freez(publish->topic);
                publish->topic = NULL;
                nd_log(NDLS_DAEMON, NDLP_ERR, "Error parsing PUBLISH message");
                return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            }
            publish->data_len = parser->mqtt_fixed_hdr_remaining_length - parser->mqtt_parsed_len;
            if (!publish->data_len) {
                publish->data = NULL;
                return MQTT_NG_CLIENT_PARSE_DONE; // 0 length payload is OK [MQTT-3.3.3]
            }
            BUF_READ_CHECK_AT_LEAST(parser->received_data, publish->data_len)

            publish->data = mallocz(publish->data_len);
            rbuf_pop(parser->received_data, publish->data, publish->data_len);
            parser->mqtt_parsed_len += publish->data_len;

            return MQTT_NG_CLIENT_PARSE_DONE;
        default:
            nd_log(NDLS_DAEMON, NDLP_ERR, "invalid state for publish varhdr parser");
            return MQTT_NG_CLIENT_INTERNAL_ERROR;
    }
    return MQTT_NG_CLIENT_OK_CALL_AGAIN;
}

// TODO move to separate file, dont send whole client pointer just to be able
// to access LOG context send parser only which should include log
static int parse_data(struct mqtt_ng_client *client)
{
    int rc;
    struct mqtt_ng_parser *parser = &client->parser;
    switch(parser->state) {
        case MQTT_PARSE_FIXED_HEADER_PACKET_TYPE:
            BUF_READ_CHECK_AT_LEAST(parser->received_data, 1)
            rbuf_pop(parser->received_data, (char*)&parser->mqtt_control_packet_type, 1);
            vbi_parser_reset_ctx(&parser->vbi_parser);
            parser->state = MQTT_PARSE_FIXED_HEADER_LEN;
            break;
        case MQTT_PARSE_FIXED_HEADER_LEN:
            rc = vbi_parser_parse(&parser->vbi_parser, parser->received_data);
            if (rc == MQTT_NG_CLIENT_PARSE_DONE) {
                parser->mqtt_fixed_hdr_remaining_length = parser->vbi_parser.result;
                parser->state = MQTT_PARSE_VARIABLE_HEADER;
                parser->varhdr_state = MQTT_PARSE_VARHDR_INITIAL;
                break;
            }
            return rc;
        case MQTT_PARSE_VARIABLE_HEADER:
            switch (get_control_packet_type(parser->mqtt_control_packet_type)) {
                case MQTT_CPT_CONNACK:
                    rc = parse_connack_varhdr(client);
                    if (rc == MQTT_NG_CLIENT_PARSE_DONE) {
                        parser->state = MQTT_PARSE_MQTT_PACKET_DONE;
                        break;
                    }
                    return rc;
                case MQTT_CPT_PUBACK:
                    rc = parse_puback_varhdr(client);
                    if (rc == MQTT_NG_CLIENT_PARSE_DONE) {
                        parser->state = MQTT_PARSE_MQTT_PACKET_DONE;
                        break;
                    }
                    return rc;
                case MQTT_CPT_SUBACK:
                    rc = parse_suback_varhdr(client);
                    if (rc != MQTT_NG_CLIENT_NEED_MORE_BYTES && rc != MQTT_NG_CLIENT_OK_CALL_AGAIN) {
                        freez(parser->mqtt_packet.suback.reason_codes);
                    }
                    if (rc == MQTT_NG_CLIENT_PARSE_DONE) {
                        parser->state = MQTT_PARSE_MQTT_PACKET_DONE;
                        break;
                    }
                    return rc;
                case MQTT_CPT_PUBLISH:
                    rc = parse_publish_varhdr(client);
                    if (rc == MQTT_NG_CLIENT_PARSE_DONE) {
                        parser->state = MQTT_PARSE_MQTT_PACKET_DONE;
                        break;
                    }
                    return rc;
                case MQTT_CPT_PINGRESP:
                    if (parser->mqtt_fixed_hdr_remaining_length) {
                        nd_log(NDLS_DAEMON, NDLP_ERR, "PINGRESP has to be 0 Remaining Length."); // [MQTT-3.13.1]
                        return MQTT_NG_CLIENT_PROTOCOL_ERROR;
                    }
                    parser->state = MQTT_PARSE_MQTT_PACKET_DONE;
                    ping_timeout = 0;
                    break;
                case MQTT_CPT_DISCONNECT:
                    rc = parse_disconnect_varhdr(client);
                    if (rc == MQTT_NG_CLIENT_PARSE_DONE) {
                        parser->state = MQTT_PARSE_MQTT_PACKET_DONE;
                        break;
                    }
                    return rc;
                default:
                    nd_log(NDLS_DAEMON, NDLP_ERR, "Parsing Control Packet Type %" PRIu8 " not implemented yet.", get_control_packet_type(parser->mqtt_control_packet_type));
                    rbuf_bump_tail(parser->received_data, parser->mqtt_fixed_hdr_remaining_length);
                    parser->state = MQTT_PARSE_MQTT_PACKET_DONE;
                    return MQTT_NG_CLIENT_NOT_IMPL_YET;
            }
            // we could also return MQTT_NG_CLIENT_OK_CALL_AGAIN
            // and be called again later
            /* FALLTHROUGH */
        case MQTT_PARSE_MQTT_PACKET_DONE:
            parser->state = MQTT_PARSE_FIXED_HEADER_PACKET_TYPE;
            return MQTT_NG_CLIENT_MQTT_PACKET_DONE;
    }
    return MQTT_NG_CLIENT_OK_CALL_AGAIN;
}

// set next MQTT fragment to send
// return 1 if nothing to send
// return -1 on error
// return 0 if there is fragment set
static int mqtt_ng_next_to_send(struct mqtt_ng_client *client) {
    if (client->client_state == MQTT_STATE_CONNECT_PENDING) {
        client->main_buffer.sending_frag = client->connect_msg;
        client->client_state = MQTT_STATE_CONNECTING;
        return 0;
    }
    if (client->client_state != MQTT_STATE_CONNECTED)
        return -1;

    struct buffer_fragment *frag = BUFFER_FIRST_FRAG(&client->main_buffer.hdr_buffer);
    while (frag) {
        // Skip fragments marked for garbage collection - their data may have
        // been freed by mark_message_for_gc() after a timeout or ACK
        if (frag_is_marked_for_gc(frag)) {
            frag = frag->next;
            continue;
        }
        if (frag->sent != frag->len)
            break;
        frag = frag->next;
    }

    if ( client->ping_pending && (!frag || (frag->flags & BUFFER_FRAG_MQTT_PACKET_HEAD && frag->sent == 0)) ) {
        client->ping_pending = 0;
        client->ping_frag.sent = 0;
        client->ping_frag.sent_monotonic_ut = 0;
        client->main_buffer.sending_frag = &client->ping_frag;
        return 0;
    }

    client->main_buffer.sending_frag = frag;
    return frag == NULL ? 1 : 0;
}

// send current fragment
// return 0 if whole remaining length could be sent as a whole
// return -1 if send buffer was filled and
// nothing could be written anymore
// return 1 if last fragment of a message was fully sent
static int send_fragment(struct mqtt_ng_client *client) {
    worker_is_busy(WORKER_ACLK_SEND_FRAGMENT);

    struct buffer_fragment *frag = client->main_buffer.sending_frag;

    // for readability
    unsigned char *ptr = frag->data + frag->sent;
    size_t bytes = frag->len - frag->sent;

    size_t processed = 0;

    if (bytes)
        processed = client->send_fnc_ptr(client->user_ctx, ptr, bytes);
    else
        nd_log(NDLS_DAEMON, NDLP_WARNING, "This fragment was fully sent already. This should not happen!");

    frag->sent_monotonic_ut = now_monotonic_usec();
    frag->sent += processed;
    if (frag->sent != frag->len)
        return -1;

    if (frag->flags & BUFFER_FRAG_MQTT_PACKET_TAIL) {
        client->time_of_last_send = time(NULL);
        if (frag != &client->ping_frag)
            __atomic_fetch_sub(&client->stats.tx_messages_queued, 1, __ATOMIC_RELAXED);
        __atomic_fetch_add(&client->stats.tx_messages_sent, 1, __ATOMIC_RELAXED);
        client->main_buffer.sending_frag = NULL;
        return 1;
    }

    client->main_buffer.sending_frag = frag->next;
    
    return 0;
}

// attempt sending all fragments of current single MQTT packet
static int send_all_message_fragments(struct mqtt_ng_client *client) {
    int rc;
    while ( !(rc = send_fragment(client)) );
    return rc;
}

static void try_send_all(struct mqtt_ng_client *client) {
    do {
        if (client->main_buffer.sending_frag == NULL && mqtt_ng_next_to_send(client))
            return;
    } while(send_all_message_fragments(client) >= 0);
}

int handle_incoming_traffic(struct mqtt_ng_client *client)
{
    int rc;
    while ((rc = parse_data(client)) == MQTT_NG_CLIENT_OK_CALL_AGAIN) {
        ;
    }
    if (rc != MQTT_NG_CLIENT_MQTT_PACKET_DONE)
        return rc;

    struct mqtt_publish *pub;
    struct mqtt_property *prop;
    __atomic_fetch_add(&client->stats.rx_messages_rcvd, 1, __ATOMIC_RELAXED);

    uint8_t ctrl_packet_type = get_control_packet_type(client->parser.mqtt_control_packet_type);
    switch (ctrl_packet_type) {
        case MQTT_CPT_CONNACK:
            worker_is_busy(WORKER_ACLK_CPT_CONNACK);

            LOCK_HDR_BUFFER(&client->main_buffer);
            // Invalidate sending_frag if it points to any fragment in the CONNECT message
            if (sending_frag_in_message(&client->main_buffer, client->connect_msg))
                client->main_buffer.sending_frag = NULL;
            mark_message_for_gc(client->connect_msg);
            UNLOCK_HDR_BUFFER(&client->main_buffer);

            client->connect_msg = NULL;

            if (client->client_state != MQTT_STATE_CONNECTING) {
                nd_log(NDLS_DAEMON, NDLP_ERR, "ACLK: Received unexpected CONNACK");
                client->client_state = MQTT_STATE_ERROR;
                return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            }

            if ((prop = get_property_by_id(client->parser.properties_parser.head, MQTT_PROP_MAX_PKT_SIZE)) != NULL) {
                nd_log(NDLS_DAEMON, NDLP_INFO, "ACLK: MQTT server limits message size to %" PRIu32, prop->data.uint32);
                client->max_msg_size = prop->data.uint32;
            }

            if ((prop = get_property_by_id(client->parser.properties_parser.head, MQTT_PROP_RECEIVE_MAX)) != NULL) {
                nd_log(NDLS_DAEMON, NDLP_INFO, "ACLK: MQTT server receive maximum is %" PRIu16, prop->data.uint16);
                __atomic_store_n(&client->rx_maximum, prop->data.uint16, __ATOMIC_RELAXED);
            }

            if (client->connack_callback)
                client->connack_callback(client->user_ctx, client->parser.mqtt_packet.connack.reason_code);
            if (!client->parser.mqtt_packet.connack.reason_code) {
                nd_log(NDLS_DAEMON, NDLP_INFO, "ACLK: MQTT Connection Accepted By Server");
                client->client_state = MQTT_STATE_CONNECTED;
                break;
            }
            client->client_state = MQTT_STATE_ERROR;
            return MQTT_NG_CLIENT_SERVER_RETURNED_ERROR;

        case MQTT_CPT_PUBACK:
            worker_is_busy(WORKER_ACLK_CPT_PUBACK);

            if (mark_packet_acked(client, client->parser.mqtt_packet.puback.packet_id))
                return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            if (client->puback_callback)
                client->puback_callback(client->parser.mqtt_packet.puback.packet_id);
            break;

        case MQTT_CPT_PINGRESP:
            worker_is_busy(WORKER_ACLK_CPT_PINGRESP);
            usec_t latency = now_monotonic_usec() - client->ping_frag.sent_monotonic_ut;
            pulse_aclk_sent_message_acked(latency, client->ping_frag.len);
            break;

        case MQTT_CPT_SUBACK:
            worker_is_busy(WORKER_ACLK_CPT_SUBACK);
            if (mark_packet_acked(client, client->parser.mqtt_packet.suback.packet_id))
                return MQTT_NG_CLIENT_PROTOCOL_ERROR;
            break;

        case MQTT_CPT_PUBLISH:
            worker_is_busy(WORKER_ACLK_CPT_PUBLISH);
            pub = &client->parser.mqtt_packet.publish;

            if (pub->qos > 1) {
                freez(pub->topic);
                freez(pub->data);
                return MQTT_NG_CLIENT_NOT_IMPL_YET;
            }

            if ( pub->qos == 1 && ((rc = mqtt_ng_puback(client, pub->packet_id, 0))) ) {
                client->client_state = MQTT_STATE_ERROR;
                nd_log(NDLS_DAEMON, NDLP_ERR, "Error generating PUBACK reply for PUBLISH");
                return rc;
            }

            if ( (prop = get_property_by_id(client->parser.properties_parser.head, MQTT_PROP_TOPIC_ALIAS)) != NULL ) {
                // Topic Alias property was sent from server
                void *topic_ptr;
                uint16_t topic_alias = prop->data.uint16;
                if (!c_rhash_get_ptr_by_uint64(client->rx_aliases, topic_alias, &topic_ptr)) {
                    if (pub->topic != NULL) {
                        nd_log(NDLS_DAEMON, NDLP_ERR, "We do not yet support topic alias reassignment");
                        return MQTT_NG_CLIENT_NOT_IMPL_YET;
                    }
                    pub->topic = topic_ptr;
                } else {
                    if (pub->topic == NULL) {
                        nd_log(NDLS_DAEMON, NDLP_ERR, "Topic alias with id %" PRIu16 " unknown and topic not set by server!", topic_alias);
                        return MQTT_NG_CLIENT_PROTOCOL_ERROR;
                    }
                    c_rhash_insert_uint64_ptr(client->rx_aliases, topic_alias, pub->topic);
                }
            }

            if (client->msg_callback) {
                worker_is_busy(WORKER_ACLK_MSG_CALLBACK);
                client->msg_callback(pub->topic, pub->data, pub->data_len, pub->qos);
            }

            // in case we have property topic alias and we have topic we take over the string
            // and add pointer to it into topic alias list
            if (prop == NULL)
                freez(pub->topic);
            freez(pub->data);
            return MQTT_NG_CLIENT_WANT_WRITE;

        case MQTT_CPT_DISCONNECT:
            worker_is_busy(WORKER_ACLK_CPT_DISCONNECT);
            nd_log(NDLS_DAEMON, NDLP_INFO, "Got MQTT DISCONNECT control packet from server. Reason code: %d", (int)client->parser.mqtt_packet.disconnect.reason_code);
            client->client_state = MQTT_STATE_DISCONNECTED;
            break;

        default:
            worker_is_busy(WORKER_ACLK_CPT_UNKNOWN);
            nd_log(NDLS_DAEMON, NDLP_INFO, "Got unknown control packet %u from server", ctrl_packet_type);
            break;
    }

    return rc;
}

#define PACKET_TIMEOUT_REPEAT_CHECK (60)

int mqtt_ng_sync(struct mqtt_ng_client *client)
{
    if (client->client_state == MQTT_STATE_RAW || client->client_state == MQTT_STATE_DISCONNECTED)
        return 0;
    
    if (client->client_state == MQTT_STATE_ERROR)
        return 1;

    // Check for packet timeouts and cleanup
    static time_t last_maintenance = 0;
    if (now_realtime_sec() - last_maintenance >= PACKET_TIMEOUT_REPEAT_CHECK) {
        // if check packet returns true then we did max cleanup, possibly there are more packets to cleanup
        // so do not update last_maintenance thus forcing check again
        if (likely(!check_packet_monitor_list_for_timeouts(client)))
            last_maintenance = now_realtime_sec();
    }

    worker_is_busy(WORKER_ACLK_TRY_SEND_ALL);

    LOCK_HDR_BUFFER(&client->main_buffer);
    try_send_all(client);
    UNLOCK_HDR_BUFFER(&client->main_buffer);

    int rc;

    worker_is_busy(WORKER_ACLK_HANDLE_INCOMING);
    while ((rc = handle_incoming_traffic(client)) != MQTT_NG_CLIENT_NEED_MORE_BYTES) {
        if (rc < 0)
            break;
        if (rc == MQTT_NG_CLIENT_WANT_WRITE) {
            worker_is_busy(WORKER_ACLK_TRY_SEND_ALL);

            LOCK_HDR_BUFFER(&client->main_buffer);
            try_send_all(client);
            UNLOCK_HDR_BUFFER(&client->main_buffer);

            worker_is_busy(WORKER_ACLK_HANDLE_INCOMING);
        }
    }

    if (rc < 0)
        return rc;

    return 0;
}

static int mqtt_ng_unittest_push_bytes(rbuf_t buffer, const char *data, size_t len)
{
    return rbuf_push(buffer, data, len) == len ? 0 : 1;
}

#define MQTT_NG_TEST(condition, msg) do {                                      \
        if (!(condition)) {                                                    \
            fprintf(stderr, "mqtt_ng unittest FAILED: %s (%s:%d)\n",          \
                    (msg), __FUNCTION__, __LINE__);                            \
            errors++;                                                          \
        }                                                                      \
    } while(0)

static struct {
    const char *topic;
    char payload[16];
    size_t payload_len;
    int qos;
    unsigned calls;
} mqtt_ng_unittest_msg;

static void mqtt_ng_unittest_msg_callback(const char *topic, const void *msg, size_t msglen, int qos)
{
    mqtt_ng_unittest_msg.topic = topic;
    mqtt_ng_unittest_msg.payload_len = msglen;
    mqtt_ng_unittest_msg.qos = qos;
    mqtt_ng_unittest_msg.calls++;

    size_t copy_len = MIN(msglen, sizeof(mqtt_ng_unittest_msg.payload) - 1);
    if (copy_len)
        memcpy(mqtt_ng_unittest_msg.payload, msg, copy_len);
    mqtt_ng_unittest_msg.payload[copy_len] = '\0';
}

static int mqtt_ng_unittest_reject_malformed_properties_packet(
    const char *packet_name,
    const char *packet,
    size_t packet_len,
    size_t expected_unread)
{
    int errors = 0;
    rbuf_t input = rbuf_create(128, 128);
    struct mqtt_ng_init settings = {
        .data_in = input,
        .data_out_fnc = NULL,
        .user_ctx = NULL,
        .connack_callback = NULL,
        .puback_callback = NULL,
        .msg_callback = NULL,
    };
    struct mqtt_ng_client *client = mqtt_ng_init(&settings);

    MQTT_NG_TEST(client != NULL, packet_name);
    if (!client) {
        rbuf_free(input);
        return errors;
    }

    MQTT_NG_TEST(!mqtt_ng_unittest_push_bytes(input, packet, packet_len), packet_name);

    int rc = handle_incoming_traffic(client);
    MQTT_NG_TEST(rc == MQTT_NG_CLIENT_PROTOCOL_ERROR, packet_name);
    MQTT_NG_TEST(rbuf_bytes_available(input) == expected_unread, packet_name);

    mqtt_ng_destroy(client);
    rbuf_free(input);
    return errors;
}

static int mqtt_ng_unittest_2byte_field_lengths(void)
{
    int errors = 0;
    size_t too_long_len = (size_t)UINT16_MAX + 1;
    char *field = mallocz(too_long_len + 1);
    memset(field, 'x', too_long_len);
    field[too_long_len] = '\0';

    uint16_t encoded_len = 0;
    field[UINT16_MAX] = '\0';
    MQTT_NG_TEST(
        mqtt_ng_get_2byte_field_length("test field", field, &encoded_len) && encoded_len == UINT16_MAX,
        "two-byte field accepts exactly 65535 bytes");
    field[UINT16_MAX] = 'x';
    MQTT_NG_TEST(
        !mqtt_ng_get_2byte_field_length("test field", field, &encoded_len),
        "two-byte field rejects 65536 bytes");

    rbuf_t input = rbuf_create(128, 128);
    struct mqtt_ng_init settings = {
        .data_in = input,
        .data_out_fnc = NULL,
        .user_ctx = NULL,
        .connack_callback = NULL,
        .puback_callback = NULL,
        .msg_callback = NULL,
    };
    struct mqtt_ng_client *client = mqtt_ng_init(&settings);
    MQTT_NG_TEST(client != NULL, "mqtt_ng_init succeeds for two-byte field test");
    if (!client) {
        rbuf_free(input);
        freez(field);
        return errors;
    }

    struct mqtt_auth_properties auth = {.client_id = field};
    MQTT_NG_TEST(
        mqtt_ng_generate_connect(&client->main_buffer, &auth, NULL, 60) == NULL,
        "CONNECT rejects oversized Client ID");

    auth.client_id = "client";
    auth.username = field;
    MQTT_NG_TEST(
        mqtt_ng_generate_connect(&client->main_buffer, &auth, NULL, 60) == NULL,
        "CONNECT rejects oversized User Name");

    auth.username = NULL;
    auth.password = field;
    MQTT_NG_TEST(
        mqtt_ng_generate_connect(&client->main_buffer, &auth, NULL, 60) == NULL,
        "CONNECT rejects oversized Password");

    auth.password = NULL;
    struct mqtt_lwt_properties lwt = {.will_topic = field};
    MQTT_NG_TEST(
        mqtt_ng_generate_connect(&client->main_buffer, &auth, &lwt, 60) == NULL,
        "CONNECT rejects oversized Will Topic");

    struct mqtt_sub sub = {.topic = field};
    MQTT_NG_TEST(
        mqtt_ng_subscribe(client, &sub, 1) == MQTT_NG_MSGGEN_USER_ERROR,
        "SUBSCRIBE rejects oversized Topic Filter");

    char payload = 'x';
    uint16_t packet_id = 0;
    MQTT_NG_TEST(
        mqtt_ng_publish(
            client, field, NULL, &payload, CALLER_RESPONSIBILITY, sizeof(payload), 0, &packet_id) ==
            MQTT_NG_MSGGEN_USER_ERROR,
        "PUBLISH rejects oversized Topic Name");

    MQTT_NG_TEST(
        BUFFER_BYTES_USED(&client->main_buffer.hdr_buffer) == 0,
        "oversized two-byte fields queue no fragments");

    mqtt_ng_destroy(client);
    rbuf_free(input);
    freez(field);
    return errors;
}

static int mqtt_ng_unittest_topic_alias_uint16(void)
{
    int errors = 0;
    rbuf_t input = rbuf_create(128, 128);
    struct mqtt_ng_init settings = {
        .data_in = input,
        .data_out_fnc = NULL,
        .user_ctx = NULL,
        .connack_callback = NULL,
        .puback_callback = NULL,
        .msg_callback = mqtt_ng_unittest_msg_callback,
    };
    struct mqtt_ng_client *client = mqtt_ng_init(&settings);

    const char publish_alias_44[] = {
        (char)(MQTT_CPT_PUBLISH << 4), 16,
        0, 7, 'a', 'l', 'i', 'a', 's', '4', '4',
        3, MQTT_PROP_TOPIC_ALIAS, 0x00, 0x2c,
        'l', 'o', 'w',
    };
    const char publish_alias_256[] = {
        (char)(MQTT_CPT_PUBLISH << 4), 17,
        0, 8, 'a', 'l', 'i', 'a', 's', '2', '5', '6',
        3, MQTT_PROP_TOPIC_ALIAS, 0x01, 0x00,
        'm', 'i', 'd',
    };
    const char publish_alias_300[] = {
        (char)(MQTT_CPT_PUBLISH << 4), 17,
        0, 8, 'a', 'l', 'i', 'a', 's', '3', '0', '0',
        3, MQTT_PROP_TOPIC_ALIAS, 0x01, 0x2c,
        'b', 'i', 'g',
    };
    const char publish_alias_only[] = {
        (char)(MQTT_CPT_PUBLISH << 4), 9,
        0, 0,
        3, MQTT_PROP_TOPIC_ALIAS, 0x01, 0x2c,
        't', 'w', 'o',
    };

    memset(&mqtt_ng_unittest_msg, 0, sizeof(mqtt_ng_unittest_msg));

    MQTT_NG_TEST(client != NULL, "mqtt_ng_init succeeds for uint16 topic alias test");
    if (!client) {
        rbuf_free(input);
        return errors;
    }

    MQTT_NG_TEST(!mqtt_ng_unittest_push_bytes(input, publish_alias_44, sizeof(publish_alias_44)),
                 "push PUBLISH with alias 44");

    int rc = handle_incoming_traffic(client);
    MQTT_NG_TEST(rc == MQTT_NG_CLIENT_WANT_WRITE, "PUBLISH with alias 44 parses");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.calls == 1, "PUBLISH with alias 44 invokes callback");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.topic && !strcmp(mqtt_ng_unittest_msg.topic, "alias44"),
                 "PUBLISH with alias 44 callback topic");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.payload_len == 3 && !strcmp(mqtt_ng_unittest_msg.payload, "low"),
                 "PUBLISH with alias 44 callback payload");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.qos == 0, "PUBLISH with alias 44 callback qos");

    MQTT_NG_TEST(!mqtt_ng_unittest_push_bytes(input, publish_alias_256, sizeof(publish_alias_256)),
                 "push PUBLISH with alias 256");

    rc = handle_incoming_traffic(client);
    MQTT_NG_TEST(rc == MQTT_NG_CLIENT_WANT_WRITE, "PUBLISH with alias 256 parses");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.calls == 2, "PUBLISH with alias 256 invokes callback");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.topic && !strcmp(mqtt_ng_unittest_msg.topic, "alias256"),
                 "PUBLISH with alias 256 callback topic");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.payload_len == 3 && !strcmp(mqtt_ng_unittest_msg.payload, "mid"),
                 "PUBLISH with alias 256 callback payload");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.qos == 0, "PUBLISH with alias 256 callback qos");

    MQTT_NG_TEST(!mqtt_ng_unittest_push_bytes(input, publish_alias_300, sizeof(publish_alias_300)),
                 "push PUBLISH with alias 300");

    rc = handle_incoming_traffic(client);
    MQTT_NG_TEST(rc == MQTT_NG_CLIENT_WANT_WRITE, "PUBLISH with alias 300 parses");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.calls == 3, "PUBLISH with alias 300 invokes callback");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.topic && !strcmp(mqtt_ng_unittest_msg.topic, "alias300"),
                 "PUBLISH with alias 300 callback topic");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.payload_len == 3 && !strcmp(mqtt_ng_unittest_msg.payload, "big"),
                 "PUBLISH with alias 300 callback payload");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.qos == 0, "PUBLISH with alias 300 callback qos");

    MQTT_NG_TEST(!mqtt_ng_unittest_push_bytes(input, publish_alias_only, sizeof(publish_alias_only)),
                 "push alias-only PUBLISH with alias 300");

    rc = handle_incoming_traffic(client);
    MQTT_NG_TEST(rc == MQTT_NG_CLIENT_WANT_WRITE, "alias-only PUBLISH with alias 300 parses");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.calls == 4, "alias-only PUBLISH with alias 300 invokes callback");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.topic && !strcmp(mqtt_ng_unittest_msg.topic, "alias300"),
                 "alias-only PUBLISH with alias 300 callback topic");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.payload_len == 3 && !strcmp(mqtt_ng_unittest_msg.payload, "two"),
                 "alias-only PUBLISH with alias 300 callback payload");
    MQTT_NG_TEST(mqtt_ng_unittest_msg.qos == 0, "alias-only PUBLISH with alias 300 callback qos");

    mqtt_ng_destroy(client);
    rbuf_free(input);
    return errors;
}

static int mqtt_ng_unittest_reset_after_partial_publish(void)
{
    int errors = 0;
    rbuf_t input = rbuf_create(128, 128);
    struct mqtt_ng_init settings = {
        .data_in = input,
        .data_out_fnc = NULL,
        .user_ctx = NULL,
        .connack_callback = NULL,
        .puback_callback = NULL,
        .msg_callback = NULL,
    };
    struct mqtt_ng_client *client = mqtt_ng_init(&settings);

    const char partial_publish[] = {
        (char)(MQTT_CPT_PUBLISH << 4), 23,
        0, 4, 't', 'e', 's', 't',
        0,
    };

    MQTT_NG_TEST(client != NULL, "mqtt_ng_init succeeds");
    if (!client) {
        rbuf_free(input);
        return errors;
    }

    MQTT_NG_TEST(!mqtt_ng_unittest_push_bytes(input, partial_publish, sizeof(partial_publish)),
                 "push partial PUBLISH");

    int rc = handle_incoming_traffic(client);
    MQTT_NG_TEST(rc == MQTT_NG_CLIENT_NEED_MORE_BYTES, "partial PUBLISH needs more bytes");
    MQTT_NG_TEST(client->parser.state == MQTT_PARSE_VARIABLE_HEADER, "partial PUBLISH stays in variable-header parser");
    MQTT_NG_TEST(client->parser.varhdr_state == MQTT_PARSE_PAYLOAD, "partial PUBLISH waits for payload");
    MQTT_NG_TEST(client->parser.mqtt_packet.publish.topic != NULL, "partial PUBLISH owns parsed topic");

    struct mqtt_auth_properties auth = {
        .client_id = "unit-test",
    };

    MQTT_NG_TEST(mqtt_ng_connect(client, &auth, NULL, 60) == 0, "mqtt_ng_connect resets parser");
    MQTT_NG_TEST(client->parser.state == MQTT_PARSE_FIXED_HEADER_PACKET_TYPE, "parser state reset");
    MQTT_NG_TEST(client->parser.mqtt_packet.publish.topic == NULL, "partial PUBLISH topic released");
    MQTT_NG_TEST(client->parser.properties_parser.head == NULL, "partial properties released");

    const char normal_publish[] = {
        (char)(MQTT_CPT_PUBLISH << 4), 7,
        0, 2, 'o', 'k',
        0,
        'h', 'i',
    };

    MQTT_NG_TEST(!mqtt_ng_unittest_push_bytes(input, normal_publish, sizeof(normal_publish)),
                 "push normal PUBLISH after reconnect");

    rc = handle_incoming_traffic(client);
    MQTT_NG_TEST(rc == MQTT_NG_CLIENT_WANT_WRITE, "normal PUBLISH parses after reconnect reset");
    MQTT_NG_TEST(client->parser.state == MQTT_PARSE_FIXED_HEADER_PACKET_TYPE, "parser ready after normal PUBLISH");

    mqtt_ng_destroy(client);
    rbuf_free(input);
    return errors;
}

static int mqtt_ng_unittest_partial_suback_reason_codes(void)
{
    int errors = 0;
    rbuf_t input = rbuf_create(128, 128);
    struct mqtt_ng_init settings = {
        .data_in = input,
        .data_out_fnc = NULL,
        .user_ctx = NULL,
        .connack_callback = NULL,
        .puback_callback = NULL,
        .msg_callback = NULL,
    };
    struct mqtt_ng_client *client = mqtt_ng_init(&settings);

    const char partial_suback[] = {
        (char)(MQTT_CPT_SUBACK << 4), 5,
        0, 1,
        0,
        0,
    };
    const char remaining_suback[] = {
        (char)0x80,
    };

    MQTT_NG_TEST(client != NULL, "mqtt_ng_init succeeds for SUBACK parser test");
    if (!client) {
        rbuf_free(input);
        return errors;
    }

    MQTT_NG_TEST(!mqtt_ng_unittest_push_bytes(input, partial_suback, sizeof(partial_suback)),
                 "push partial SUBACK");

    int rc = handle_incoming_traffic(client);
    MQTT_NG_TEST(rc == MQTT_NG_CLIENT_NEED_MORE_BYTES, "partial SUBACK needs more bytes");
    MQTT_NG_TEST(client->parser.state == MQTT_PARSE_VARIABLE_HEADER, "partial SUBACK stays in variable-header parser");
    MQTT_NG_TEST(client->parser.varhdr_state == MQTT_PARSE_REASONCODES, "partial SUBACK waits for reason codes");
    MQTT_NG_TEST(client->parser.mqtt_packet.suback.reason_code_count == 2, "partial SUBACK reason-code count");
    MQTT_NG_TEST(client->parser.mqtt_packet.suback.reason_codes_pending == 1, "partial SUBACK pending reason code");
    uint8_t *reason_codes = client->parser.mqtt_packet.suback.reason_codes;
    MQTT_NG_TEST(reason_codes != NULL, "partial SUBACK owns reason-code buffer");
    if (reason_codes)
        MQTT_NG_TEST(reason_codes[0] == 0, "partial SUBACK stores first reason code");

    MQTT_NG_TEST(!mqtt_ng_unittest_push_bytes(input, remaining_suback, sizeof(remaining_suback)),
                 "push remaining SUBACK reason code");

    rc = parse_suback_varhdr(client);
    MQTT_NG_TEST(rc == MQTT_NG_CLIENT_PARSE_DONE, "fragmented SUBACK reason codes parse done");
    if (reason_codes) {
        MQTT_NG_TEST(reason_codes[0] == 0, "fragmented SUBACK preserves first reason code");
        MQTT_NG_TEST(reason_codes[1] == 0x80, "fragmented SUBACK appends second reason code");
    }
    MQTT_NG_TEST(client->parser.mqtt_packet.suback.reason_codes_pending == 0, "fragmented SUBACK consumes all reason codes");

    freez(client->parser.mqtt_packet.suback.reason_codes);
    client->parser.mqtt_packet.suback.reason_codes = NULL;
    mqtt_ng_destroy(client);
    rbuf_free(input);
    return errors;
}

static int mqtt_ng_unittest_malformed_suback_properties_length(void)
{
    int errors = 0;
    rbuf_t input = rbuf_create(128, 128);
    struct mqtt_ng_init settings = {
        .data_in = input,
        .data_out_fnc = NULL,
        .user_ctx = NULL,
        .connack_callback = NULL,
        .puback_callback = NULL,
        .msg_callback = NULL,
    };
    struct mqtt_ng_client *client = mqtt_ng_init(&settings);

    const char malformed_suback[] = {
        (char)(MQTT_CPT_SUBACK << 4), 4,
        0, 1,
        2,
        MQTT_PROP_REQ_PROBLEM_INFO,
        (char)(MQTT_CPT_PINGRESP << 4), 0,
    };

    MQTT_NG_TEST(client != NULL, "mqtt_ng_init succeeds for malformed SUBACK parser test");
    if (!client) {
        rbuf_free(input);
        return errors;
    }

    MQTT_NG_TEST(!mqtt_ng_unittest_push_bytes(input, malformed_suback, sizeof(malformed_suback)),
                 "push malformed SUBACK");

    int rc = handle_incoming_traffic(client);
    MQTT_NG_TEST(rc == MQTT_NG_CLIENT_PROTOCOL_ERROR, "malformed SUBACK properties length is rejected");
    MQTT_NG_TEST(client->parser.mqtt_packet.suback.reason_codes == NULL,
                 "malformed SUBACK does not allocate reason codes");
    MQTT_NG_TEST(rbuf_bytes_available(input) == 3,
                 "malformed SUBACK leaves following bytes unread");

    mqtt_ng_destroy(client);
    rbuf_free(input);
    return errors;
}

static int mqtt_ng_unittest_malformed_ack_properties_length(void)
{
    int errors = 0;
    const char malformed_connack[] = {
        (char)(MQTT_CPT_CONNACK << 4), 4,
        0, 0,
        1,
        MQTT_PROP_REASON_STR,
        (char)(MQTT_CPT_PINGRESP << 4), 0,
    };
    const char malformed_disconnect[] = {
        (char)(MQTT_CPT_DISCONNECT << 4), 3,
        0,
        1,
        MQTT_PROP_REASON_STR,
        (char)(MQTT_CPT_PINGRESP << 4), 0,
    };
    const char malformed_puback[] = {
        (char)(MQTT_CPT_PUBACK << 4), 5,
        0, 1,
        0,
        1,
        MQTT_PROP_REASON_STR,
        (char)(MQTT_CPT_PINGRESP << 4), 0,
    };

    errors += mqtt_ng_unittest_reject_malformed_properties_packet(
        "malformed CONNACK properties do not consume following packet",
        malformed_connack,
        sizeof(malformed_connack),
        2);
    errors += mqtt_ng_unittest_reject_malformed_properties_packet(
        "malformed DISCONNECT properties do not consume following packet",
        malformed_disconnect,
        sizeof(malformed_disconnect),
        2);
    errors += mqtt_ng_unittest_reject_malformed_properties_packet(
        "malformed PUBACK properties do not consume following packet",
        malformed_puback,
        sizeof(malformed_puback),
        2);

    return errors;
}

int mqtt_ng_unittest(void)
{
    int errors = 0;

    fprintf(stderr, "\nrunning mqtt_ng unittest\n");

    errors += mqtt_ng_unittest_2byte_field_lengths();
    errors += mqtt_ng_unittest_reset_after_partial_publish();
    errors += mqtt_ng_unittest_topic_alias_uint16();
    errors += mqtt_ng_unittest_partial_suback_reason_codes();
    errors += mqtt_ng_unittest_malformed_suback_properties_length();
    errors += mqtt_ng_unittest_malformed_ack_properties_length();

    if (errors)
        fprintf(stderr, "mqtt_ng unittest: %d ERROR(S)\n", errors);
    else
        fprintf(stderr, "mqtt_ng unittest: OK\n");

    return errors;
}

time_t mqtt_ng_last_send_time(struct mqtt_ng_client *client)
{
    return client->time_of_last_send;
}

void mqtt_ng_set_max_mem(struct mqtt_ng_client *client, size_t bytes)
{
    client->max_mem_bytes = bytes;
}

void mqtt_ng_get_stats(struct mqtt_ng_client *client, struct mqtt_ng_stats *stats)
{
    stats->tx_messages_queued = __atomic_load_n(&client->stats.tx_messages_queued, __ATOMIC_RELAXED);
    stats->tx_messages_sent = __atomic_load_n(&client->stats.tx_messages_sent, __ATOMIC_RELAXED);
    stats->rx_messages_rcvd = __atomic_load_n(&client->stats.rx_messages_rcvd, __ATOMIC_RELAXED);
    stats->packets_waiting_puback = __atomic_load_n(&client->stats.packets_waiting_puback, __ATOMIC_RELAXED);
    stats->rx_maximum = __atomic_load_n(&client->rx_maximum, __ATOMIC_RELAXED);

    stats->tx_bytes_queued = 0;
    stats->tx_buffer_reclaimable = 0;
    stats->max_puback_wait_us = 0;
    stats->max_send_queue_wait_us = 0;
    stats->max_unsent_wait_us = 0;
    stats->max_partial_wait_us = 0;

    // First pass: compute buffer usage/queued bytes and max send-queue wait time (unsent messages)
    LOCK_HDR_BUFFER(&client->main_buffer);
    stats->tx_buffer_used = BUFFER_BYTES_USED(&client->main_buffer.hdr_buffer);
    stats->tx_buffer_free = BUFFER_BYTES_AVAILABLE(&client->main_buffer.hdr_buffer);
    stats->tx_buffer_size = client->main_buffer.hdr_buffer.size;
    struct buffer_fragment *frag = BUFFER_FIRST_FRAG(&client->main_buffer.hdr_buffer);
    usec_t now_ut = now_monotonic_usec();
    while (frag) {
        stats->tx_bytes_queued += frag->len - frag->sent;
        if (frag_is_marked_for_gc(frag))
            stats->tx_buffer_reclaimable += FRAG_SIZE_IN_BUFFER(frag);

        // For HEAD fragments that are not fully sent yet (unsent or partially sent),
        // track max send-queue wait time. Prefer the enqueue timestamp; if missing, fall back to first-send.
        if ((frag->flags & BUFFER_FRAG_MQTT_PACKET_HEAD) && frag->sent < frag->len) {
            usec_t base = frag->enqueued_monotonic_ut ? frag->enqueued_monotonic_ut : frag->sent_monotonic_ut;
            if (base) {
                usec_t waited = now_ut - base;
                uint64_t w = (uint64_t)waited;
                if (frag->sent == 0) {
                    if (w > stats->max_unsent_wait_us) stats->max_unsent_wait_us = w;
                } else {
                    if (w > stats->max_partial_wait_us) stats->max_partial_wait_us = w;
                }
                if (w > stats->max_send_queue_wait_us) stats->max_send_queue_wait_us = w;
            } else {
                // Throttled debug if enqueue time is missing on an unsent HEAD
                if (frag->sent == 0) {
                    static time_t last_warn = 0;
                    time_t now_s = now_monotonic_sec();
                    if (now_s - last_warn > 60) {
                        nd_log(NDLS_DAEMON, NDLP_DEBUG, "ACLK: Missing enqueue timestamp on unsent MQTT packet head");
                        last_warn = now_s;
                    }
                }
            }
        }

        frag = frag->next;
    }
    UNLOCK_HDR_BUFFER(&client->main_buffer);

    // Second pass: compute max PUBACK wait time by correlating HEAD fragments with pending packet IDs
    spinlock_lock(&client->pending_packets.spinlock);
    LOCK_HDR_BUFFER(&client->main_buffer);
    frag = BUFFER_FIRST_FRAG(&client->main_buffer.hdr_buffer);
    while (frag) {
        if ((frag->flags & BUFFER_FRAG_MQTT_PACKET_HEAD) && frag->packet_id) {
            Pvoid_t *Pvalue = JudyLGet(client->pending_packets.JudyL, (Word_t)frag->packet_id, PJE0);
            if (Pvalue) {
                // message is still pending PUBACK
                if (frag->sent_monotonic_ut) {
                    usec_t waited = now_ut - frag->sent_monotonic_ut;
                    if ((uint64_t)waited > stats->max_puback_wait_us)
                        stats->max_puback_wait_us = (uint64_t)waited;
                }
            }
        }
        frag = frag->next;
    }
    UNLOCK_HDR_BUFFER(&client->main_buffer);
    spinlock_unlock(&client->pending_packets.spinlock);
}

int mqtt_ng_set_topic_alias(struct mqtt_ng_client *client, const char *topic)
{
    uint16_t idx;
    spinlock_lock(&client->tx_topic_aliases.spinlock);

    if (client->tx_topic_aliases.idx_assigned >= client->tx_topic_aliases.idx_max) {
        spinlock_unlock(&client->tx_topic_aliases.spinlock);
        nd_log(NDLS_DAEMON, NDLP_ERR, "Tx topic alias indexes were exhausted (current version of the library doesn't support reassigning yet. Feel free to contribute.");
        return 0; //0 is not a valid topic alias
    }

    struct topic_alias_data *alias;
    if (!c_rhash_get_ptr_by_str(client->tx_topic_aliases.stoi_dict, topic, (void**)&alias)) {
        // this is not a problem for library but might be helpful to warn user
        // as it might indicate bug in their program (but also might be expected)
        idx = alias->idx;
        spinlock_unlock(&client->tx_topic_aliases.spinlock);
        nd_log(NDLS_DAEMON, NDLP_DEBUG, "%s topic \"%s\" already has alias set. Ignoring.", __FUNCTION__, topic);
        return idx;
    }

    alias = mallocz(sizeof(struct topic_alias_data));
    idx = ++client->tx_topic_aliases.idx_assigned;
    alias->idx = idx;
    __atomic_store_n(&alias->usage_count, 0, __ATOMIC_SEQ_CST);

    c_rhash_insert_str_ptr(client->tx_topic_aliases.stoi_dict, topic, (void*)alias);

    spinlock_unlock(&client->tx_topic_aliases.spinlock);
    return idx;
}
