/*
 * Networking binary buffers pack/unpack library
 *
 * (c) Alexander Vdolainen 2016 <avdolainen@zoho.com>
 * (c) Alexander Vdolainen 2017 <alex@vapaa.xyz>
 *
 * libtbusd is free software: you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation, either version 2.1 of the License, or
 * (at your option) any later version.
 *
 * libtbusd is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 * See the GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.";
 *
 */

#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include <sys/types.h>
#include <arpa/inet.h>
#include <errno.h>

#include <ndbuf/ndbuf.h>

static inline uint64_t ntohll(uint64_t n)
{
#ifdef WORDS_BIGENDIAN
  return n;
#else /* WORDS_BIGENDIAN */
  return (uint64_t)(((uint64_t)(n) << 56) |                      \
         (((uint64_t)(n) << 40) & 0xff000000000000ULL) | \
         (((uint64_t)(n) << 24) & 0xff0000000000ULL) | \
         (((uint64_t)(n) << 8)  & 0xff00000000ULL) | \
         (((uint64_t)(n) >> 8)  & 0xff000000ULL) | \
         (((uint64_t)(n) >> 24) & 0xff0000ULL) | \
         (((uint64_t)(n) >> 40) & 0xff00ULL) | \
         ((uint64_t)(n)  >> 56));
#endif /* WORDS_BIGENDIAN */
}

/* allocation, freeing */
#define DEFAULT_PREALLOC_SIZE  256
#define NDBUF_MAXLENGTH  6*(1024*1024)

/* allocate raw buffer with defaults preallocation */
ndbuf_t *ndbuf_new(void)
{
  ndbuf_t *b = malloc(sizeof(ndbuf_t));

  if(!b) return NULL;

  if(!(b->raw = malloc(DEFAULT_PREALLOC_SIZE))) {
    free(b);
    b = NULL;
  } else {
    b->rlength = DEFAULT_PREALLOC_SIZE;
    b->ulength = b->curr = 0;
    b->flags = 0;
  }

  return b;
}

/* create ndbuf from raw buffer */
ndbuf_t *ndbuf_new_frombuf(char *buf, size_t buf_len, void (*freebuf)(char *))
{
  ndbuf_t *b = NULL;

  if(!buf || !buf_len) return NULL;
  if(!(b = malloc(sizeof(ndbuf_t)))) return NULL;

  /* init buffer */
  b->raw = buf;
  b->rlength = b->ulength = buf_len;
  b->flags = b->curr = 0;
  b->freebuf = freebuf;
  b->flags |= NDBUF_NREA;

  return b;
}

/* will do the same as ndbuf_new but will allocate given length */
ndbuf_t *ndbuf_new_palloc(uint32_t alen)
{
  ndbuf_t *b = NULL;

  if(alen > NDBUF_MAXLENGTH) return NULL;

  if(!(b = malloc(sizeof(ndbuf_t)))) return NULL;
  else memset(b, 0, sizeof(ndbuf_t));

  if(!(b->raw = malloc(alen))) {
    free(b);
    b = NULL;
  } else b->rlength = alen;

  return b;
}

/* free all allocated space and buffer itself */
void ndbuf_free(ndbuf_t *b)
{
  if(!b) return;

  if((b->flags & NDBUF_BURN) && b->raw)
    memset(b->raw, 0, b->rlength);

  if((b->flags & NDBUF_NREA) && b->freebuf) b->freebuf(b->raw);
  else if(b->raw) free(b->raw);

  if(b->flags & NDBUF_BURN) memset(b, 0, sizeof(ndbuf_t));
  free(b);

  return;
}

/* read/write */
/* read different types, should return the size of the
 * ridden data, otherwise error occurs */
uint32_t ndbuf_read_u8(ndbuf_t *b, uint8_t *o)
{
  uint8_t u;

  if(!b || !o) return 0;
  if(!b->raw) return 0;

  if(b->curr >= b->ulength) return 0;
  u = *((uint8_t*)(b->raw + b->curr));
  b->curr += sizeof(uint8_t);
  *o = u;

  return sizeof(uint8_t);
}

uint32_t ndbuf_read_u16(ndbuf_t *b, uint16_t *o)
{
  uint16_t u;

  if(!b || !o) return 0;
  if(!b->raw) return 0;

  if(b->curr >= b->ulength) return 0;
  if(b->curr + sizeof(uint16_t) > b->ulength) return 0;

  u = *((uint16_t *)(b->raw + b->curr));
  b->curr += sizeof(uint16_t);
  *o = ntohs(u);

  return sizeof(uint16_t);
}

uint32_t ndbuf_read_u32(ndbuf_t *b, uint32_t *o)
{
  uint32_t u;

  if(!b || !o) return 0;
  if(!b->raw) return 0;

  if(b->curr >= b->ulength) return 0;
  if(b->curr + sizeof(uint32_t) > b->ulength) return 0;

  u = *((uint32_t *)(b->raw + b->curr));
  b->curr += sizeof(uint32_t);
  *o = ntohl(u);

  return sizeof(uint32_t);
}

uint32_t ndbuf_read_u64(ndbuf_t *b, uint64_t *o)
{
  uint64_t u;

  if(!b || !o) return 0;
  if(!b->raw) return 0;

  if(b->curr >= b->ulength) return 0;
  if(b->curr + sizeof(uint64_t) > b->ulength) return 0;

  u = *((uint64_t *)(b->raw + b->curr));
  b->curr += sizeof(uint64_t);
  *o = ntohll(u);

  return sizeof(uint64_t);
}

/* get raw data, pointer must be allocated with at least required length,
 * will return size of ridden data
 */
uint32_t ndbuf_read_raw(ndbuf_t *b, void *o, uint32_t rlen)
{
  if(!b || !o) return 0;
  if(!rlen || !b->raw) return 0;
  if(b->curr >= b->ulength) return 0;
  if(b->curr + rlen > b->ulength) return 0;

  memcpy(o, (const void*)b->raw + b->curr, rlen);
  b->curr += rlen;

  return rlen;
}

static int __rdb_grow(ndbuf_t *b, uint32_t len)
{
  uint32_t rlen;
  char *ne = NULL;

  if(!len) return -1;
  if(b->rlength + len > NDBUF_MAXLENGTH) return -1;

  rlen = len +
    (len%DEFAULT_PREALLOC_SIZE != 0 ? (DEFAULT_PREALLOC_SIZE - len%DEFAULT_PREALLOC_SIZE) : 0);
  if(b->rlength + rlen > NDBUF_MAXLENGTH) rlen = len;

  rlen += b->rlength;
  if(!(ne = malloc(rlen))) return -1;

  if(b->flags & NDBUF_BURN) memset(ne, 0, rlen);

  memcpy(ne, b->raw, b->ulength);

  if(b->flags & NDBUF_BURN) memset(b->raw, 0, b->rlength);
  free(b->raw);
  b->raw = ne;
  b->rlength = rlen;

  return 0;
}

/* write different types, should return the size of the
 * written data, otherwise error occurs */
uint32_t ndbuf_write_u8(ndbuf_t *b, uint8_t u)
{
  if(!b || !b->raw) return 0;
  if(b->ulength == b->rlength) {
    if(__rdb_grow(b, sizeof(uint8_t))) return 0;
  }

  *(uint8_t *)(b->raw + b->ulength) = u;
  b->ulength += sizeof(uint8_t);

  return sizeof(uint8_t);
}

uint32_t ndbuf_write_u16(ndbuf_t *b, uint16_t uu)
{
  if(!b || !b->raw) return 0;
  if(b->ulength + sizeof(uint16_t) >= b->rlength) {
    if(__rdb_grow(b, sizeof(uint16_t))) return 0;
  }

  *(uint16_t *)(b->raw + b->ulength) = htons(uu);
  b->ulength += sizeof(uint16_t);

  return sizeof(uint16_t);
}

uint32_t ndbuf_write_u32(ndbuf_t *b, uint32_t uu)
{
  if(!b || !b->raw) return 0;
  if(b->ulength + sizeof(uint32_t) >= b->rlength) {
    if(__rdb_grow(b, sizeof(uint32_t))) return 0;
  }

  *(uint32_t *)(b->raw + b->ulength) = htonl(uu);
  b->ulength += sizeof(uint32_t);

  return sizeof(uint32_t);
}

uint32_t ndbuf_write_u64(ndbuf_t *b, uint64_t uu)
{
  if(!b || !b->raw) return 0;
  if(b->ulength + sizeof(uint64_t) >= b->rlength) {
    if(__rdb_grow(b, sizeof(uint64_t))) return 0;
  }

  *(uint64_t *)(b->raw + b->ulength) = ntohll(uu);
  b->ulength += sizeof(uint64_t);

  return sizeof(uint64_t);
}

/* write raw data with the given length */
uint32_t ndbuf_write_raw(ndbuf_t *b, void *wi, uint32_t len)
{
  if(!b || !b->raw) return 0;
  if(!wi || !len) return 0;
  if(b->ulength + len >= b->rlength) {
    if(__rdb_grow(b, len)) return 0;
  }

  memcpy((void *)b->raw + b->ulength, (const void *)wi, len);
  b->ulength += len;

  return len;
}

/* write raw data *before* existing data */
uint32_t ndbuf_write_raw_head(ndbuf_t *b, void *wi, uint32_t len)
{
  char *ne;
  uint32_t rlen;

  if(!b || !b->raw) return 0;
  if(!wi || !len) return 0;
  if(b->ulength + len > b->rlength) {
    /* allocate a new one and copy it right */
    if(b->rlength + len > NDBUF_MAXLENGTH) return -1;

    rlen = len +
      (len%DEFAULT_PREALLOC_SIZE != 0 ? (DEFAULT_PREALLOC_SIZE - len%DEFAULT_PREALLOC_SIZE) : 0);
    if(b->rlength + rlen > NDBUF_MAXLENGTH) rlen = len;

    rlen += b->rlength;
    if(!(ne = malloc(rlen))) return -1;

    if(b->flags & NDBUF_BURN) memset(ne, 0, rlen);

    memcpy((void *)ne + len, b->raw, b->ulength);

    if(b->flags & NDBUF_BURN) memset(b->raw, 0, b->rlength);
    free(b->raw);
    b->raw = ne;
    b->rlength = rlen;
  } else {
    memmove((void *)b->raw + b->ulength, b->raw, b->ulength);
  }
  memcpy(b->raw, wi, len);
  b->ulength += len;

  return len;
}

/* parse */
int ndbuf_escan_va(ndbuf_t *b, const char *fmt, int argc, va_list ap)
{
  va_list ap_copy;
  union {
    uint8_t *_u8;
    uint16_t *_u16;
    uint32_t *_u32;
    uint64_t *_u64;
    void **_dp;
    char **_cstr;
    ndbuf_t **_rdb;
  } d;
  const char *t, *last;
  uint32_t len, clen;
  int r, count;

  va_copy(ap_copy, ap);

  for(t = fmt, count = 0; *t != '\0'; t++, count++) {
    if(count > argc && argc != -1) {
      va_end(ap_copy);
      return EINVAL;
    }

    switch(*t) {
    case 'b':
      d._u8 = va_arg(ap, uint8_t *);
      len = ndbuf_read_u8(b, d._u8);
      r = (len == sizeof(uint8_t)) ? 0 : -1;
      break;
    case 'w':
      d._u16 = va_arg(ap, uint16_t *);
      len = ndbuf_read_u16(b, d._u16);
      r = (len == sizeof(uint16_t)) ? 0 : -1;
      break;
    case 'd':
      d._u32 = va_arg(ap, uint32_t *);
      len = ndbuf_read_u32(b, d._u32);
      r = (len == sizeof(uint32_t)) ? 0 : -1;
      break;
    case 'q':
      d._u64 = va_arg(ap, uint64_t *);
      len = ndbuf_read_u64(b, d._u64);
      r = (len == sizeof(uint64_t)) ? 0 : -1;
      break;
    case 's':
      d._cstr = va_arg(ap, char **);
      *d._cstr = NULL;
      len = ndbuf_read_u32(b, &clen);
      if(len != sizeof(uint32_t)) {
      __errrbread:
        r = -1;
        break;
      }
      if(clen > NDBUF_MAXLENGTH) goto __errrbread;
      if((*d._cstr = malloc(clen + sizeof(char))) == NULL) {
        r = -ENOMEM;
        break;
      }
      len = ndbuf_read_raw(b, *d._cstr, clen);
      if(len != clen) {
        free(*d._cstr);
        goto __errrbread;
      }
      (*d._cstr)[len] = '\0';
      d._cstr = NULL;
      r = 0;
      break;
    case 'p':
      clen = va_arg(ap, size_t);

      d._dp = va_arg(ap, void **);
      *d._dp = NULL;
      count++;

      if((*d._dp = malloc(clen)) == NULL) {
        r = -ENOMEM;
        break;
      }
      len = ndbuf_read_raw(b, *d._dp, clen);
      if(len != clen) {
        free(*d._dp);
        goto __errrbread;
      }
      d._dp = NULL;
      r = 0;
      break;
    case 'P':
      d._dp = va_arg(ap, void **);
      *d._dp = NULL;

      len = ndbuf_read_u32(b, &clen);
      if(len != sizeof(uint32_t)) {
        r = -1;
        break;
      }
      if((*d._dp = malloc(clen)) == NULL) {
        r = -ENOMEM;
        break;
      }
      len = ndbuf_read_raw(b, *d._dp, clen);
      if(len != clen) {
        free(*d._dp);
        goto __errrbread;
      }
      d._dp = NULL;
      r = 0;
      break;
    case 'R':
      d._rdb = va_arg(ap, ndbuf_t **);
      *d._rdb = NULL;

      len = ndbuf_read_u32(b, &clen);
      if(len != sizeof(uint32_t)) {
        r = -1;
        break;
      }
      /* allocate rdb and write to it */
      if((*d._rdb = ndbuf_new_palloc(clen)) == NULL) {
        r = -ENOMEM;
        break;
      } else ndbuf_setflags(*d._rdb, NDBUF_BURN);
      len = ndbuf_read_raw(b, ndbuf_rdata(*d._rdb), clen);
      if(len != clen) {
        ndbuf_free(*d._dp);
        goto __errrbread;
      } else ndbuf_setlength(*d._rdb, clen);

      d._rdb = NULL;
      r = 0;
      break;
    }

    if(r) break;
  }

  if(argc != -1 && argc != count) r = -EINVAL;

  if(!r) {
    clen = va_arg(ap, uint32_t);
    if(clen != NDBUF_TERMINAT) {
      if(argc == -1) r = -EINVAL;
      else abort();
    }
  }

  /* clean up all the stuff on error */
  if(r) {
    last = t;
    for(t = fmt; t < last; ++t) {
      switch(*t) {
      case 'b':
      case 'w':
      case 'd':
      case 'q':
        (void)va_arg(ap_copy, void *);
        break;
      case 's':
        d._cstr = va_arg(ap_copy, char **);
        if(*d._cstr) {
          memset(*d._cstr, 0, strlen(*d._cstr));
          free(*d._cstr);
        }
        break;
      case 'R':
        d._rdb = va_arg(ap_copy, ndbuf_t **);
        if(*d._rdb) ndbuf_free(*d._rdb);
        break;
      case 'p':
        (void)va_arg(ap_copy, size_t);
      case 'P':
        d._dp = va_arg(ap_copy, void **);
        if(*d._dp) free(*d._dp);
        break;
      default:
        (void)va_arg(ap_copy, void *);
        break;
      }
    }
  }
  va_end(ap_copy);

  return r;
}

int ndbuf_escan_wot(ndbuf_t *b, const char *fmt, int argc, ...)
{
  va_list ap;
  int r;

  va_start(ap, argc);
  r = ndbuf_escan_va(b, fmt, argc, ap);
  va_end(ap);

  return r;
}

/* print FIXME: TODO: add errno setting upon failure */
uint32_t ndbuf_print_va(ndbuf_t *b, const char *fmt, int argc, va_list ap)
{
  const char *t;
  char *cstr;
  ndbuf_t *rdb;
  union {
    uint8_t _u8;
    uint16_t _u16;
    uint32_t _u32;
    uint64_t _u64;
    void *_dp;
  } d;
  uint32_t len = 0, clen;
  int r, count;

  for(t = fmt, count = 0; *t != '\0'; t++, count++) {
    if(count > argc && argc != -1) return 0;
    switch(*t) {
    case 'b':
      d._u8 = (uint8_t)va_arg(ap, unsigned int);
      r = ndbuf_write_u8(b, d._u8);
      if(r != sizeof(uint8_t)) return 0;
      else len += r;
      break;
    case 'w':
      d._u16 = (uint16_t)va_arg(ap, unsigned int);
      r = ndbuf_write_u16(b, d._u16);
      if(r != sizeof(uint16_t)) return 0;
      else len += r;
      break;
    case 'd':
      d._u32 = (uint32_t)va_arg(ap, uint32_t);
      r = ndbuf_write_u32(b, d._u32);
      if(r != sizeof(uint32_t)) return 0;
      else len += r;
      break;
    case 'q':
      d._u64 = (uint64_t)va_arg(ap, uint64_t);
      r = ndbuf_write_u64(b, d._u64);
      if(r != sizeof(uint64_t)) return 0;
      else len += r;
      break;
    case 's':
      cstr = (char *)va_arg(ap, char *);
      clen = strlen(cstr);
      r = ndbuf_write_u32(b, (uint32_t)clen);
      if(r != sizeof(uint32_t)) return 0;
      else len += r;
      r = ndbuf_write_raw(b, cstr, clen);
      if(r != clen) return 0;
      else len += clen;
      break;
    case 'R':
      rdb = (ndbuf_t *)va_arg(ap, ndbuf_t *);
      clen = ndbuf_length(rdb);
      r = ndbuf_write_u32(b, (uint32_t)clen);
      if(r != sizeof(uint32_t)) return 0;
      else len += r;
      r = ndbuf_write_raw(b, ndbuf_rdata(rdb), (uint32_t)clen);
      if(r != clen) return 0;
      else len += clen;
      break;
    case 'P':
      clen = va_arg(ap, size_t);
      d._dp = va_arg(ap, void *);
      count++;
      r = ndbuf_write_u32(b, (uint32_t)clen);
      if(r != sizeof(uint32_t)) return 0;
      else len += r;
      r = ndbuf_write_raw(b, d._dp, (uint32_t)clen);
      if(r != clen) return 0;
      else len += clen;
      break;
    case 'p':
      clen = va_arg(ap, size_t);
      d._dp = va_arg(ap, void *);
      count++;
      r = ndbuf_write_raw(b, d._dp, (uint32_t)clen);
      if(r != clen) return 0;
      else len += clen;
      break;
    default:
      return 0;
      break;
    }
  }

  if(argc != count && argc != -1) return 0;

  clen = va_arg(ap, uint32_t);
  if(clen != NDBUF_TERMINAT) {
    if(argc == -1) return 0;
    else abort(); /* corruption */
  }

  return len;
}

uint32_t ndbuf_print_wot(ndbuf_t *b, const char *fmt, int argc, ...)
{
  va_list ap;
  uint32_t wrl;

  va_start(ap, argc);
  wrl = ndbuf_print_va(b, fmt, argc, ap);
  va_end(ap);

  return wrl;
}

/* misc */
/* returns length of used space in the buffer */
uint32_t ndbuf_length(ndbuf_t *b)
{
  if(!b || !b->raw) return 0;

  return b->ulength;
}

/* returns length of allocated space in the buffer */
uint32_t ndbuf_alength(ndbuf_t *b)
{
  if(!b || !b->raw) return 0;

  return b->rlength;
}

/* returns length of the left data to read */
uint32_t ndbuf_leftlength(ndbuf_t *b)
{
  if(!b || !b->raw) return 0;
  if(b->curr > b->ulength) return 0;

  return b->ulength - b->curr;
}

/* reaturn a pointer to the raw data */
void *ndbuf_rdata(ndbuf_t *b)
{
  if(!b || !b->raw) return NULL;

  return (void *)b->raw;
}

/* return pointer to the data currently being read/write */
void *ndbuf_rdatacur(ndbuf_t *b)
{
  if(!b || !b->raw) return NULL;

  if(b->curr >= b->ulength) return NULL;

  return (void *)(b->raw + b->curr);
}

/* reset current pointer */
void ndbuf_resetcur(ndbuf_t *b)
{
  if(!b || b->raw) return;
  else b->curr = 0;

  return;
}

/* set flags to the raw buffer */
void ndbuf_setflags(ndbuf_t *b, int af)
{
  if(!b) return;

  b->flags |= af;

  return;
}

/* exchange flags for raw buff */
void ndbuf_exflags(ndbuf_t *b, int nf)
{
  if(!b) return;

  b->flags = nf;

  return;
}

int ndbuf_setlength(ndbuf_t *b, uint32_t s)
{
  if(!b) return EINVAL;

  if(s > b->rlength) return EINVAL;
  else b->ulength = s;

  return 0;
}

/* compare buffers, returns 0 if equal */
int ndbuf_cmp(ndbuf_t *b1, ndbuf_t *b2)
{
  int c = 0;

  if(!b1 | !b2) return -1;

  if((c = ndbuf_length(b1) - ndbuf_length(b2)) != 0) return c;
  else c = memcmp((const void *)ndbuf_rdata(b1),
                  (const void *)ndbuf_rdata(b2),
                  ndbuf_length(b1));

  return c;
}


/* let the buffer to use actually used bytes, not all allocated space
 * will return 0 on success (or in case if it doesn't required),
 * ENOMEM or other error if fails
 */
int ndbuf_memopt(ndbuf_t *b)
{
  uint32_t len;
  char *ne;

  if(!b || !b->raw) return EINVAL;

  if((b->rlength - b->ulength) > DEFAULT_PREALLOC_SIZE) {
    len = b->ulength +
      (b->ulength%DEFAULT_PREALLOC_SIZE != 0 ?
       (DEFAULT_PREALLOC_SIZE - b->ulength%DEFAULT_PREALLOC_SIZE) : 0);
    if(!(ne = malloc(len))) return ENOMEM;

    if(b->flags & NDBUF_BURN) memset(ne, 0, len);
    memcpy(ne, b->raw, b->ulength);
    if(b->flags & NDBUF_BURN) memset(b->raw, 0, b->ulength);
    free(b->raw);
    b->raw = ne;
    b->rlength = len;
  }

  return 0;
}

#undef DEFAULT_PREALLOC_SIZE