/*
 * netq.h
 *
 * Copyright (C) 2022 bzt (bztsrc@gitlab) MIT license
 *
 * Permission is hereby granted, free of charge, to any person
 * obtaining a copy of this software and associated documentation
 * files (the "Software"), to deal in the Software without
 * restriction, including without limitation the rights to use, copy,
 * modify, merge, publish, distribute, sublicense, and/or sell copies
 * of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 *
 * @brief Simple reliable datagram implementation with circular FIFO network queue
 * https://gitlab.com/bztsrc/netq
 */

#ifndef NETQ_H
#define NETQ_H

#ifdef __cplusplus
extern "C" {
#endif

#include <stdint.h>
#include <string.h>

/* the `nq` is a netq_t type, opaque to the application.
 * the `net` is a raw network layer type, as you please, opaque to NetQ. */

int  netq_send(void *nq, const void *msg, int len, void *net);  /* send a normal, reliable message with data */
int  netq_rst (void *nq, void *net);                            /* send a non-reliable, non-message RESET packet */
int  netq_ack (void *nq, void *net);                            /* send a non-reliable, non-message ACK packet */
int  netq_nack(void *nq, void *net);                            /* send a non-reliable, non-message NACK packet */

int  netq_push(void *nq, const void *raw, int len, void *net);  /* push a received raw packet to the queue */
int  netq_pop (void *nq, void *msg, int len);                   /* pop a message in correct order from the queue */
int  netq_pend(void *nq);                                       /* check if there's any message pending */

void netq_dump(void *nq);                                       /* dump queues to stdout (for debugging) */

#if defined(NETQ_CONFIG) || defined(NETQ_IMPLEMENTATION)
#ifndef NETQ_HASCONFIG
#define NETQ_HASCONFIG

/*** configuration ***/
#ifndef NETQ_MTU
#define NETQ_MTU        1280
#endif
#if NETQ_MTU < 576 || NETQ_MTU > 1500
#error "Unlikely NETQ_MTU value. Use something between 576 and 1500"
#endif

#ifndef NETQ_SEQ_BITS
#define NETQ_SEQ_BITS   16
#endif

#ifndef NETQ_ACK_BITS
#define NETQ_ACK_BITS   16
#endif

#ifndef NETQ_NACK_TRES
#define NETQ_NACK_TRES  24
#endif

#ifndef NETQ_PACKED
#ifndef _MSC_VER
#define NETQ_PACKED     __attribute__((packed))
#else
#define NETQ_PACKED
#endif
#endif

#if (defined(NETQ_MUTEX_TYPE) && (!defined(NETQ_MUTEX_LOCK) || !defined(NETQ_MUTEX_UNLOCK))) || \
    (defined(NETQ_MUTEX_LOCK) && (!defined(NETQ_MUTEX_TYPE) || !defined(NETQ_MUTEX_UNLOCK))) || \
    (defined(NETQ_MUTEX_UNLOCK) && (!defined(NETQ_MUTEX_TYPE) || !defined(NETQ_MUTEX_LOCK)))
#error "All of NETQ_MUTEX_TYPE, NETQ_MUTEX_LOCK and NETQ_MUTEX_UNLOCK must be defined together."
#endif
#if !defined(NETQ_MUTEX_TYPE) && !defined(NETQ_NO_PTHREAD) && defined(_PTHREAD_H)
#define NETQ_MUTEX_TYPE         pthread_mutex_t
#define NETQ_MUTEX_LOCK(m)      pthread_mutex_lock(m)
#define NETQ_MUTEX_UNLOCK(m)    pthread_mutex_unlock(m)
#endif

/*** calculated from configuration ***/
#if NETQ_SEQ_BITS == 8
#define NETQ_SEQ uint8_t
#define NETQ_SEQ_MAX 64             /* calculate as 2^(NETQ_SEQ_BITS-2) */
#define NETQ_SEQ_MSK 0x7f           /* calculate as (2^NETQ_SEQ_BITS)-1 */
#define NETQ_HTON
#define NETQ_NTOH
#elif NETQ_SEQ_BITS == 16
#define NETQ_SEQ uint16_t
#define NETQ_SEQ_MAX 16384
#define NETQ_SEQ_MSK 0x7fff
#define NETQ_HTON htons
#define NETQ_NTOH ntohs
#elif NETQ_SEQ_BITS == 32
#define NETQ_SEQ uint32_t
#define NETQ_SEQ_MAX 1073741824
#define NETQ_SEQ_MSK 0x7fffffff
#define NETQ_HTON htonl
#define NETQ_NTOH ntohl
#else
#error "Invalid NETQ_SEQ_BITS, can be 8, 16 or 32"
#endif

#if NETQ_ACK_BITS == 8
#define NETQ_ACK uint8_t
#elif NETQ_ACK_BITS == 16
#define NETQ_ACK uint16_t
#elif NETQ_ACK_BITS == 32
#define NETQ_ACK uint32_t
#elif NETQ_ACK_BITS == 64
#define NETQ_ACK uint64_t
#else
#error "Invalid NETQ_ACK_BITS, can be 8, 16, 32 or 64"
#endif
#ifndef NETQ_QUEUE_SIZE
#define NETQ_QUEUE_SIZE NETQ_ACK_BITS
#endif
#if NETQ_QUEUE_SIZE < NETQ_ACK_BITS
#error "NETQ_QUEUE_SIZE must be at least NETQ_ACK_BITS"
#endif
#if NETQ_NACK_TRES < 0 || NETQ_NACK_TRES >= 2 * NETQ_ACK_BITS
#error "Unlikely NETQ_NACK_TRES value. Use something like 1.5 * NETQ_ACK_BITS"
#endif

/* RFC 1982 compatible sequence number comparator */
#define NETQ_SEQ_EQ(i1,i2)  ((i1)==(i2))
#define NETQ_SEQ_LE(i1,i2)  ((i1)<(i2) && ((i2)-(i1)<NETQ_SEQ_MAX))

/* the message header */
enum { NETQ_MSG_TYPE_MSG, NETQ_MSG_TYPE_RST, NETQ_MSG_TYPE_ACK, NETQ_MSG_TYPE_NACK };
#define NETQ_MSG_TYPE(m) ((((m)->seq >> 7) & 1) | (((m)->ack >> 6) & 2))

#ifdef _MSC_VER
#pragma pack(push)
#pragma pack(1)
#endif
typedef struct {
    NETQ_SEQ seq;               /* this packet's sequence number,     MSB: type bit 0 */
    NETQ_SEQ ack;               /* ack bitmask first sequence number, MSB: type bit 1 */
    NETQ_ACK msk;               /* ack bitmask */
} NETQ_PACKED netq_hdr_t;
#ifdef _MSC_VER
#pragma pack(pop)
#endif

/* the network queue context */
#define NETQ_QUE_IDX(s) ((((s) & NETQ_SEQ_MSK) & (NETQ_QUEUE_SIZE - 1)) * (NETQ_MTU + 2))  /* get position in queue */
#define NETQ_QUE_LEN(b) ((((b)[1] & 0xf) << 8) | (b)[0])                  /* size of the message in the queue */
#define NETQ_QUE_ACK(b) ((b)[1] & 0x10)                                   /* queue entry got ack */
typedef struct {
    NETQ_SEQ seq_out;           /* outgoing sequence number */
    NETQ_SEQ seq_in;            /* incoming sequence number */
    NETQ_SEQ seq_ack;           /* largest sequence number that got acked */
    NETQ_SEQ seq_pop;           /* next packet to be popped */
    uint8_t send[NETQ_QUEUE_SIZE * (NETQ_MTU + 2)]; /* the outgoing queue */
    uint8_t recv[NETQ_QUEUE_SIZE * (NETQ_MTU + 2)]; /* the incoming queue */
#ifdef NETQ_MUTEX_TYPE
    NETQ_MUTEX_TYPE mutex;
#endif
} netq_t;
#endif /* NETQ_HASCONFIG */

#endif /* NETQ_CONFIG || NETQ_IMPLEMENTATION */

#ifdef NETQ_IMPLEMENTATION

#if !defined(NETQ_SEND)
#error "You must define a NETQ_SEND function"
#endif
int NETQ_SEND(void *ctx, const void *msg, int len);

/*** Private functions ***/

void _netq_wrack(netq_t *ctx, netq_hdr_t *hdr)
{
    NETQ_SEQ i, j;
    uint8_t *buf;
#ifdef __BIG_ENDIAN
    uint8_t *msk = (uint8_t*)&hdr->msk;
#endif
    j = (ctx->seq_in - 1) & NETQ_SEQ_MSK;
    hdr->ack = NETQ_HTON(j);
    hdr->msk = 0;
    for(i = 0; i < NETQ_ACK_BITS; i++) {
        buf = &ctx->recv[NETQ_QUE_IDX(j)];
        if(((netq_hdr_t*)(buf + 2))->seq == j && NETQ_QUE_LEN(buf))
#ifdef __BIG_ENDIAN
            msk[i >> 3] |= (1 << (i & 7));
#else
            hdr->msk |= (1 << i);
#endif
        j--; j &= NETQ_SEQ_MSK;
    }
}

void _netq_resend(netq_t *ctx, NETQ_SEQ a, NETQ_ACK m, void *net)
{
    NETQ_SEQ i, j, k;
    netq_hdr_t *hdr, ack_new = { 0 };
    uint8_t *buf;

    _netq_wrack(ctx, &ack_new);
    for(i = 0, j = a; i < NETQ_ACK_BITS; i++, m >>= 1) {
        k = NETQ_QUE_IDX(j); buf = &ctx->send[k]; hdr = (netq_hdr_t*)(buf + 2);
        if((m & 1) && hdr->seq == j && NETQ_QUE_LEN(buf) && !NETQ_QUE_ACK(buf)) {
            hdr->seq = NETQ_HTON(hdr->seq);
            hdr->ack = ack_new.ack; hdr->msk = ack_new.msk;
            NETQ_SEND(net, buf + 2, NETQ_QUE_LEN(buf));
            hdr->seq = NETQ_NTOH(hdr->seq);
        }
        j--; j &= NETQ_SEQ_MSK;
    }
}

/*** Public API ***/

/**
 * Add net queue send header to message
 * @param nq: the NetQ instance (opaque to the application)
 * @param msg: the message buffer
 * @param len: length of message
 * @param net: your networking layer instance (opaque to NetQ)
 * @return negative on error, number bytes sent otherwise.
 */
int netq_send(void *nq, const void *msg, int len, void *net)
{
    int ret;
    netq_t *ctx = (netq_t*)nq;
    netq_hdr_t *hdr;
    uint8_t *buf;

    if(!ctx || !msg || len < 0 || len >= NETQ_MTU - (int)sizeof(netq_hdr_t)) return -2;
#ifdef NETQ_MUTEX_TYPE
    NETQ_MUTEX_LOCK(&ctx->mutex);
#endif
    /* automatically resend all unacked older packets */
    _netq_resend(ctx, (ctx->seq_ack - NETQ_ACK_BITS) & NETQ_SEQ_MSK, -1, net);

    /* construct and send the current message */
    buf = &ctx->send[NETQ_QUE_IDX(ctx->seq_out)]; hdr = (netq_hdr_t*)(buf + 2);
    if(msg && len > 0)
        memcpy(buf + 2 + sizeof(netq_hdr_t), msg, len);
    len += sizeof(netq_hdr_t);
    buf[0] = len & 0xff; buf[1] = (len >> 8) & 0xf;
    hdr->seq = NETQ_HTON(ctx->seq_out);
    _netq_wrack(ctx, hdr);
    ret = NETQ_SEND(net, buf + 2, len);
    hdr->seq = NETQ_NTOH(hdr->seq);
    if(ret < len) { memset(buf, 0, NETQ_MTU + 2); } else { ctx->seq_out++; ctx->seq_out &= NETQ_SEQ_MSK; }
#ifdef NETQ_MUTEX_TYPE
    NETQ_MUTEX_UNLOCK(&ctx->mutex);
#endif
    return ret;
}

/**
 * Send a non-reliable, non-message reset packet
 * @param nq: the NetQ instance (opaque to the application)
 * @param net: your networking layer instance (opaque to NetQ)
 * @return negative on error, number bytes sent otherwise.
 */
int netq_rst(void *nq, void *net)
{
    netq_t *ctx = (netq_t*)nq;
    netq_hdr_t hdr = { 0 };

    if(!ctx) return -2;
    hdr.seq = NETQ_HTON(ctx->seq_out) | 0x80;
    return NETQ_SEND(net, &hdr, sizeof(netq_hdr_t));
}

/**
 * Send a non-reliable, non-message acknowledge packet
 * @param nq: the NetQ instance (opaque to the application)
 * @param net: your networking layer instance (opaque to NetQ)
 * @return negative on error, number bytes sent otherwise.
 */
int netq_ack(void *nq, void *net)
{
    netq_t *ctx = (netq_t*)nq;
    netq_hdr_t hdr = { 0 };

    if(!ctx) return -2;
    hdr.seq = NETQ_HTON(ctx->seq_out);
    _netq_wrack(ctx, &hdr);
    hdr.ack |= 0x80;
    return NETQ_SEND(net, &hdr, sizeof(netq_hdr_t));
}

/**
 * Send a non-reliable, non-message negative acknowledge (re-transmission request) packet
 * @param nq: the NetQ instance (opaque to the application)
 * @param net: your networking layer instance (opaque to NetQ)
 * @return negative on error, number bytes sent otherwise.
 */
int netq_nack(void *nq, void *net)
{
    NETQ_SEQ i, j, k, a;
    netq_t *ctx = (netq_t*)nq;
    netq_hdr_t hdr = { 0 };
#ifdef __BIG_ENDIAN
    uint8_t *msk = (uint8_t*)&hdr.msk;
#endif

    if(!ctx) return -2;
    a = (ctx->seq_in - 1 - NETQ_ACK_BITS) & NETQ_SEQ_MSK;
    for(i = 0, j = a; i < NETQ_ACK_BITS; i++) {
        k = NETQ_QUE_IDX(j);
        if(((netq_hdr_t*)(&ctx->recv[k + 2]))->seq != j || !NETQ_QUE_LEN(&ctx->recv[k]))
#ifdef __BIG_ENDIAN
            msk[i >> 3] |= (1 << (i & 7));
#else
            hdr.msk |= (1 << i);
#endif
        j--; j &= NETQ_SEQ_MSK;
    }
    if(hdr.msk) {
        hdr.seq = NETQ_HTON(ctx->seq_out) | 0x80;
        hdr.ack = NETQ_HTON(a) | 0x80;
        return NETQ_SEND(net, &hdr, sizeof(netq_hdr_t));
    }
    return 0;
}

/**
 * Push a raw packet (message with header) to the queue and send ACK or NACK if needed
 * @param nq: the NetQ instance (opaque to the application)
 * @param raw: the raw packet (message with header) buffer
 * @param len: length of raw packet
 * @param net: your networking layer instance (opaque to NetQ)
 * @return 1 on success, 0 on error (more packet lost or arrived than the queue size).
 */
int netq_push(void *nq, const void *raw, int len, void *net)
{
    NETQ_SEQ i, j, k, msg_seq, msg_ack;
    NETQ_ACK msg_msk = 0;
    int ret = 1;
    netq_t *ctx = (netq_t*)nq;
    netq_hdr_t *hdr_in, *hdr_buf, hdr_ack;
    uint8_t *buf, msg_type;
#ifdef __BIG_ENDIAN
    uint8_t *msk;
#endif

    if(!ctx || !raw || len < (int)sizeof(netq_hdr_t) || len >= NETQ_MTU) return 0;
#ifdef NETQ_MUTEX_TYPE
    NETQ_MUTEX_LOCK(&ctx->mutex);
#endif
    hdr_in = (netq_hdr_t*)raw; msg_type = NETQ_MSG_TYPE(hdr_in);
    msg_seq = NETQ_NTOH(hdr_in->seq) & NETQ_SEQ_MSK; msg_ack = NETQ_NTOH(hdr_in->ack) & NETQ_SEQ_MSK;
    buf = &ctx->recv[NETQ_QUE_IDX(msg_seq)]; hdr_buf = (netq_hdr_t*)(buf + 2);
#ifdef __BIG_ENDIAN
    msk = (uint8_t*)&hdr_in->msk;
    for(i = 0; i < NETQ_ACK_BITS; i++)
        if(msk[i >> 3] & (1 << (i & 7))) msg_msk |= (1 << i);
#else
    msg_msk = hdr_in->msk;
#endif

    switch(msg_type) {
        case NETQ_MSG_TYPE_RST:
#ifdef NETQ_MUTEX_TYPE
            /* do not clear the mutex */
            ctx->seq_out = ctx->seq_in = ctx->seq_ack = ctx->seq_pop = 0;
            memset(ctx->send, 0, sizeof(ctx->send)); memset(ctx->recv, 0, sizeof(ctx->recv));
#else
            memset(ctx, 0, sizeof(netq_t));
#endif
        break;
        case NETQ_MSG_TYPE_NACK:
            /* for NACK, the ack mask is negated, eg. contains which packets to resend */
            _netq_resend(ctx, msg_ack, msg_msk, net);
        break;
        case NETQ_MSG_TYPE_ACK:
        case NETQ_MSG_TYPE_MSG:
            /* update ack flags in send queue from the newly received message */
            j = msg_ack;
            for(i = 0; i < NETQ_ACK_BITS; i++) {
                k = NETQ_QUE_IDX(j);
                if(((netq_hdr_t*)(&ctx->send[k + 2]))->seq == j) ctx->send[k + 1] |= 0x10;
                j--; j &= NETQ_SEQ_MSK;
            }
            /* send ACK or NACK packets when needed */
            if(NETQ_SEQ_LE((ctx->seq_ack + NETQ_NACK_TRES) & NETQ_SEQ_MSK, msg_ack)) {
                netq_nack(nq, net);
                ctx->seq_ack = msg_ack;
            } else
            if(NETQ_SEQ_EQ(ctx->seq_ack, msg_ack)) {
                _netq_wrack(ctx, &hdr_ack);
                if(hdr_in->msk != hdr_ack.msk) netq_ack(nq, net);
            }
            if(msg_type == NETQ_MSG_TYPE_ACK) break;

            /* store message in queue */
            if(NETQ_SEQ_EQ(ctx->seq_in, msg_seq) || NETQ_SEQ_LE(ctx->seq_in, msg_seq))
                ctx->seq_in = (msg_seq + 1) & NETQ_SEQ_MSK;
            else if(NETQ_SEQ_LE(msg_seq, ctx->seq_in) && hdr_buf->seq == msg_seq) break;
            len -= sizeof(netq_hdr_t); buf[0] = len & 0xff; buf[1] = (len >> 8) & 0xf;
#ifdef __BIG_ENDIAN
            hdr_buf->seq = msg_seq; hdr_buf->ack = msg_ack; hdr_buf->msk = msg_msk;
            memcpy(buf + 2 + sizeof(netq_hdr_t), (uint8_t*)raw + sizeof(netq_hdr_t), len);
#else
            memcpy(buf + 2, raw, len + sizeof(netq_hdr_t));
#endif
            /* uh oh, should never happen, unless queue was too small */
            if(NETQ_SEQ_LE(msg_seq, ctx->seq_pop) || NETQ_SEQ_LE(hdr_buf->seq, ctx->seq_pop)) ret = 0;
        break;
    }
#ifdef NETQ_MUTEX_TYPE
    NETQ_MUTEX_UNLOCK(&ctx->mutex);
#endif
    return ret;
}

/**
 * Receive message in correct order
 * @param nq: the NetQ instance (opaque to the application)
 * @param msg: output buffer (should be at least NETQ_MTU big)
 * @param len: length of output buffer
 * @return -1 on error, 0 if there was no new message, otherwise the number of bytes received and message in `msg`.
 */
int netq_pop (void *nq, void *msg, int len)
{
    int l;
    netq_t *ctx = (netq_t*)nq;
    uint8_t *buf;

    if(!ctx || !msg || len < 1 || !netq_pend(nq)) return 0;
    /* no need for locking this */
    buf = &ctx->recv[NETQ_QUE_IDX(ctx->seq_pop)];
    l = NETQ_QUE_LEN(buf);
    if(l < 1 || l > len || ((netq_hdr_t*)(buf + 2))->seq != ctx->seq_pop) return -1;
    memcpy(msg, buf + 2 + sizeof(netq_hdr_t), l);
    ctx->seq_pop++; ctx->seq_pop &= NETQ_SEQ_MSK;
    return l;
}

/**
 * Return true if there's a message pending in the queue
 */
int netq_pend(void *nq)
{
    netq_t *ctx = (netq_t*)nq;
    return ctx && NETQ_SEQ_LE(ctx->seq_pop, ctx->seq_in);
}

#ifndef NETQ_NODEBUG

/**
 * Dump queue to stdout
 */
void netq_dump(void *nq)
{
#if NETQ_QUEUE_SIZE == 8
#define NETQ_STR "2"
#elif NETQ_QUEUE_SIZE == 16
#define NETQ_STR "4"
#elif NETQ_QUEUE_SIZE == 32
#define NETQ_STR "8"
#elif NETQ_QUEUE_SIZE == 64
#define NETQ_STR "16l"
#endif
    int i, j;
    netq_t *ctx = (netq_t*)nq;
    netq_hdr_t *hdr;
    uint8_t *buf;

    printf("NetQ configuration: sequence number: uint%u_t, ack mask: uint%u_t, mtu: %u, header: %u, queue: %u, memory: %u bytes"
#ifdef NETQ_MUTEX_TYPE
        ", thread-safety"
#endif
        ".\n", NETQ_SEQ_BITS, NETQ_ACK_BITS, NETQ_MTU, (unsigned int)sizeof(netq_hdr_t), NETQ_QUEUE_SIZE, (unsigned int)sizeof(netq_t));
    if(ctx) {
        printf("\nCounters: out: %u, in: %u, latest ack: %u, next pop: %u, netq_pend()=%u\n\nSender queue:\n",
            ctx->seq_out, ctx->seq_in, ctx->seq_ack, ctx->seq_pop, netq_pend(nq));
        for(i = 0, buf = ctx->send; i < NETQ_QUEUE_SIZE; i++, buf += NETQ_MTU + 2)
            if(NETQ_QUE_LEN(buf))
                printf("%2u. len:%4u, seq:%6u, acked: %u\n", i, NETQ_QUE_LEN(buf), ((netq_hdr_t*)(buf + 2))->seq, buf[1] >> 4);
        printf("\nReceiver queue:\n");
        for(i = 0, buf = ctx->recv; i < NETQ_QUEUE_SIZE; i++, buf += NETQ_MTU + 2) {
            if(NETQ_QUE_LEN(buf)) {
                hdr = (netq_hdr_t*)(buf + 2);
                printf("%2u. len:%4u, seq:%6u, ack:%6u/%0" NETQ_STR "x (seqs", i, NETQ_QUE_LEN(buf), hdr->seq, hdr->ack, hdr->msk);
                for(j = 0; j < NETQ_QUEUE_SIZE; j++) if(hdr->msk & (1 << j)) printf(" %u", hdr->ack - j);
                printf(")\n");
            }
        }
        printf("\n");
    }
#undef NETQ_STR
}

#endif /* !NETQ_NODEBUG */

#endif /* NETQ_IMPLEMENTATION */

#ifdef __cplusplus
}
#endif

#endif /* NETQ_H */
