#include <stdio.h>
#include <string.h>

#include "empack.h"

#define EMPACK_POS_FIXINT_MAX (0x7f)
#define EMPACK_FIXMAP_MAX     (0x8f)
#define EMPACK_FIXARRAY_MAX   (0x9f)
#define EMPACK_FIXSTR_MAX     (0xbf)
#define EMPACK_NEG_FIXINT_MAX (0xff)

/** Attempt to allocate size bytes from the memory arena, NULL if out of memory */
static void *empack_allocate(empack_t *empack, uint32_t size);

/** Copy src to dst but with reversed byte order */
static void memswap(uint8_t *dst, const uint8_t *src, uint32_t len);
/** Convert a big-endian UINT16 buffer to native uint16_t */
static uint16_t buf_to_u16(const uint8_t *buf);
/** Convert a big-endian UINT32 buffer to native uint32_t */
static uint32_t buf_to_u32(const uint8_t *buf);
/** Convert a big-endian UINT64 buffer to native uint64_t */
static uint64_t buf_to_u64(const uint8_t *buf);
/** Convert a big-endian INT16 buffer to native int16_t */
static int16_t buf_to_i16(const uint8_t *buf);
/** Convert a big-endian INT32 buffer to native int32_t */
static int32_t buf_to_i32(const uint8_t *buf);
/** Convert a big-endian INT64 buffer to native int64_t */
static int64_t buf_to_i64(const uint8_t *buf);
/** Convert a big-endian FLOAT32 buffer to native float */
static float buf_to_float(const uint8_t *buf);
/** Convert a big-endian FLOAT64 buffer to native double */
static double buf_to_double(const uint8_t *buf);

void
empack_init(empack_t *empack, void *arena, uint32_t arena_sz)
{
    empack->arena = arena;
    empack->arena_sz = arena_sz;
    empack->arena_used = 0;

    memset(arena, 0, arena_sz);
}

empack_status_e
empack_parse(empack_t *empack, const void *msgpack, uint32_t len)
{
    const uint8_t *buf = (const uint8_t *)msgpack;
    uint32_t p = 0;
    empack_status_e ret;
    empack_obj_t *cur = &empack->root;
    empack_obj_t *parent = NULL;
    uint32_t tokens_needed = 1;

    while (tokens_needed > 0) {
        uint32_t len_used;

        if (parent) {
            switch (parent->type) {
            case EMPACK_TYPE_FIXARRAY:
            case EMPACK_TYPE_ARRAY16:
            case EMPACK_TYPE_ARRAY32:
            {
                empack_array_t *parr = &parent->val.array;

                if (parr->object_count == parr->object_count_real) {
                    ret = EMPACK_STATUS_ERROR;
                    goto empack_parse_return;
                }

                cur = &parr->objects[parr->object_count_real];
                cur->parent = parent;

                parr->object_count_real++;
                if (parr->object_count == parr->object_count_real) {
                    parent = parent->parent;
                }
                break;
            }
            case EMPACK_TYPE_FIXMAP:
            case EMPACK_TYPE_MAP16:
            case EMPACK_TYPE_MAP32:
            {
                empack_map_t *pmap = &parent->val.map;

                if (pmap->pair_count == pmap->pair_count_real) {
                    ret = EMPACK_STATUS_ERROR;
                    goto empack_parse_return;
                }

                if (pmap->pairs[pmap->pair_count_real].key.type == EMPACK_TYPE_ERROR) {
                    /* We haven't parsed the key yet */
                    cur = &pmap->pairs[pmap->pair_count_real].key;
                } else {
                    cur = &pmap->pairs[pmap->pair_count_real].value;
                    pmap->pair_count_real++;
                }
                cur->parent = parent;

                if (pmap->pair_count == pmap->pair_count_real) {
                    parent = parent->parent;
                }
                break;
            }
            default:
                ret = EMPACK_STATUS_ERROR;
                goto empack_parse_return;
                break;
            }
        }


        len_used = len - p;
        ret = empack_parse_one(cur, &buf[p], &len_used);
        if ((ret != EMPACK_STATUS_COMPLETE) && (ret != EMPACK_STATUS_EXTRA)) {
            goto empack_parse_return;
        }

        p += len_used;
        tokens_needed--;
        if ((tokens_needed == 0) && (ret == EMPACK_STATUS_EXTRA)) {
            goto empack_parse_return;
        }


        switch (cur->type) {
        case EMPACK_TYPE_FIXARRAY:
        case EMPACK_TYPE_ARRAY16:
        case EMPACK_TYPE_ARRAY32:
        {
            empack_array_t *parr = &cur->val.array;

            parr->objects = empack_allocate(empack, sizeof(empack_obj_t) * parr->object_count);
            if (parr->objects == NULL) {
                ret = EMPACK_STATUS_NOMEM;
                goto empack_parse_return;
            }

            tokens_needed += parr->object_count;
            parent = cur;
            break;
        }

        case EMPACK_TYPE_FIXMAP:
        case EMPACK_TYPE_MAP16:
        case EMPACK_TYPE_MAP32:
        {
            empack_map_t *pmap = &cur->val.map;

            pmap->pairs = empack_allocate(empack, sizeof(empack_map_t) * pmap->pair_count);
            if (pmap->pairs == NULL) {
                ret = EMPACK_STATUS_NOMEM;
                goto empack_parse_return;
            }

            /* Initialize the key objects to a type of error so we can tell if
             * we've already parsed them when we loop back around to parse the
             * value */
            for (uint32_t i = 0; i < pmap->pair_count; i++) {
                pmap->pairs[i].key.type = EMPACK_TYPE_ERROR;
            }
            break;
        }
        default:
            break;
        }
    }

    /* prioritize errors over incomplete status */
    if ((ret == EMPACK_STATUS_COMPLETE) && (p < len)) {
        ret = EMPACK_STATUS_WANT_MORE;
    }

empack_parse_return:
    return (ret);
}

empack_status_e
empack_parse_one(empack_obj_t *object, const void *msgpack, uint32_t *len)
{
    const uint8_t *buf = (const uint8_t *)msgpack;
    uint32_t p = 0;
    empack_status_e ret = EMPACK_STATUS_EXTRA;

    memset(object, 0, sizeof(*object));

#define EMPACK_PARSEONE_ASSERT(need) do { \
    if ((p + (need)) > (*len)) {          \
        ret = EMPACK_STATUS_ERROR;        \
        goto empack_parse_one_return;     \
    }                                     \
} while (0)

    EMPACK_PARSEONE_ASSERT(1);
    /* this casting will be validated for object types with a range */
    object->type = (empack_type_e)buf[p];
    p += 1;

    switch (object->type) {
    case EMPACK_TYPE_NIL:
    case EMPACK_TYPE_TRUE:
    case EMPACK_TYPE_FALSE:
        /* no extra value associated for these types */
        break;

    case EMPACK_TYPE_BIN8:
        EMPACK_PARSEONE_ASSERT(1);
        object->val.bin.size = buf[p];
        p += 1;
        EMPACK_PARSEONE_ASSERT(object->val.bin.size);
        object->val.bin.data = &buf[p];
        p += object->val.bin.size;
        break;

    case EMPACK_TYPE_BIN16:
        EMPACK_PARSEONE_ASSERT(2);
        object->val.bin.size = buf_to_u16(&buf[p]);
        p += 2;
        EMPACK_PARSEONE_ASSERT(object->val.bin.size);
        object->val.bin.data = &buf[p];
        p += object->val.bin.size;
        break;

    case EMPACK_TYPE_BIN32:
        EMPACK_PARSEONE_ASSERT(4);
        object->val.bin.size = buf_to_u32(&buf[p]);
        p += 4;
        EMPACK_PARSEONE_ASSERT(object->val.bin.size);
        object->val.bin.data = &buf[p];
        p += object->val.bin.size;
        break;

    case EMPACK_TYPE_EXT8:
        EMPACK_PARSEONE_ASSERT(1);
        object->val.ext.type = buf[p];
        p += 1;
        EMPACK_PARSEONE_ASSERT(1);
        object->val.ext.size = buf[p];
        p += 1;
        EMPACK_PARSEONE_ASSERT(object->val.ext.size);
        object->val.ext.data = &buf[p];
        p += object->val.ext.size;
        break;

    case EMPACK_TYPE_EXT16:
        EMPACK_PARSEONE_ASSERT(1);
        object->val.ext.type = buf[p];
        p += 1;
        EMPACK_PARSEONE_ASSERT(2);
        object->val.ext.size = buf_to_u16(&buf[p]);
        p += 2;
        EMPACK_PARSEONE_ASSERT(object->val.ext.size);
        object->val.ext.data = &buf[p];
        p += object->val.ext.size;
        break;

    case EMPACK_TYPE_EXT32:
        EMPACK_PARSEONE_ASSERT(1);
        object->val.ext.type = buf[p];
        p += 1;
        EMPACK_PARSEONE_ASSERT(4);
        object->val.ext.size = buf_to_u32(&buf[p]);
        p += 4;
        EMPACK_PARSEONE_ASSERT(object->val.ext.size);
        object->val.ext.data = &buf[p];
        p += object->val.ext.size;
        break;

    case EMPACK_TYPE_FLOAT32:
        EMPACK_PARSEONE_ASSERT(4);
        object->val.float32 = buf_to_float(&buf[p]);
        p += 4;
        break;

    case EMPACK_TYPE_FLOAT64:
        EMPACK_PARSEONE_ASSERT(8);
        object->val.float64 = buf_to_double(&buf[p]);
        p += 8;
        break;

    case EMPACK_TYPE_UINT8:
        EMPACK_PARSEONE_ASSERT(1);
        object->val.uint8 = buf[p];
        p += 1;
        break;

    case EMPACK_TYPE_UINT16:
        EMPACK_PARSEONE_ASSERT(2);
        object->val.uint16 = buf_to_u16(&buf[p]);
        p += 2;
        break;

    case EMPACK_TYPE_UINT32:
        EMPACK_PARSEONE_ASSERT(4);
        object->val.uint16 = buf_to_u32(&buf[p]);
        p += 4;
        break;

    case EMPACK_TYPE_UINT64:
        EMPACK_PARSEONE_ASSERT(8);
        object->val.uint64 = buf_to_u64(&buf[p]);
        p += 8;
        break;

    case EMPACK_TYPE_INT8:
        EMPACK_PARSEONE_ASSERT(1);
        object->val.int8 = (int8_t)buf[p];
        p += 1;
        break;

    case EMPACK_TYPE_INT16:
        EMPACK_PARSEONE_ASSERT(2);
        object->val.int16 = buf_to_i16(&buf[p]);
        p += 2;
        break;

    case EMPACK_TYPE_INT32:
        EMPACK_PARSEONE_ASSERT(4);
        object->val.int32 = buf_to_i32(&buf[p]);
        p += 4;
        break;

    case EMPACK_TYPE_INT64:
        EMPACK_PARSEONE_ASSERT(8);
        object->val.int64 = buf_to_i64(&buf[p]);
        p += 8;
        break;

    case EMPACK_TYPE_FIXEXT1:
        EMPACK_PARSEONE_ASSERT(1);
        object->val.ext.type = buf[p];
        p += 1;
        EMPACK_PARSEONE_ASSERT(1);
        object->val.ext.data = &buf[p];
        object->val.ext.size = 1;
        p += 1;
        break;

    case EMPACK_TYPE_FIXEXT2:
        EMPACK_PARSEONE_ASSERT(1);
        object->val.ext.type = buf[p];
        p += 1;
        EMPACK_PARSEONE_ASSERT(2);
        object->val.ext.data = &buf[p];
        object->val.ext.size = 2;
        p += 2;
        break;

    case EMPACK_TYPE_FIXEXT4:
        EMPACK_PARSEONE_ASSERT(1);
        object->val.ext.type = buf[p];
        p += 1;
        EMPACK_PARSEONE_ASSERT(4);
        object->val.ext.data = &buf[p];
        object->val.ext.size = 4;
        p += 4;
        break;

    case EMPACK_TYPE_FIXEXT8:
        EMPACK_PARSEONE_ASSERT(1);
        object->val.ext.type = buf[p];
        p += 1;
        EMPACK_PARSEONE_ASSERT(8);
        object->val.ext.data = &buf[p];
        object->val.ext.size = 8;
        p += 8;
        break;

    case EMPACK_TYPE_FIXEXT16:
        EMPACK_PARSEONE_ASSERT(1);
        object->val.ext.type = buf[p];
        p += 1;
        EMPACK_PARSEONE_ASSERT(16);
        object->val.ext.data = &buf[p];
        object->val.ext.size = 16;
        p += 16;
        break;

    case EMPACK_TYPE_STR8:
        EMPACK_PARSEONE_ASSERT(1);
        object->val.str.size = buf[p];
        p += 1;
        EMPACK_PARSEONE_ASSERT(object->val.str.size);
        object->val.str.data = (const char *)&buf[p];
        p += object->val.str.size;
        break;

    case EMPACK_TYPE_STR16:
        EMPACK_PARSEONE_ASSERT(2);
        object->val.str.size = buf_to_u16(&buf[p]);
        p += 2;
        EMPACK_PARSEONE_ASSERT(object->val.str.size);
        object->val.str.data = (const char *)&buf[p];
        p += object->val.str.size;
        break;

    case EMPACK_TYPE_STR32:
        EMPACK_PARSEONE_ASSERT(4);
        object->val.str.size = buf_to_u32(&buf[p]);
        p += 4;
        EMPACK_PARSEONE_ASSERT(object->val.str.size);
        object->val.str.data = (const char *)&buf[p];
        p += object->val.str.size;
        break;

    case EMPACK_TYPE_ARRAY16:
        EMPACK_PARSEONE_ASSERT(2);
        object->val.array.object_count = buf_to_u16(&buf[p]);
        p += 2;
        break;

    case EMPACK_TYPE_ARRAY32:
        EMPACK_PARSEONE_ASSERT(4);
        object->val.array.object_count = buf_to_u32(&buf[p]);
        p += 4;
        break;

    case EMPACK_TYPE_MAP16:
        EMPACK_PARSEONE_ASSERT(2);
        object->val.map.pair_count = buf_to_u16(&buf[p]);
        p += 2;
        break;

    case EMPACK_TYPE_MAP32:
        EMPACK_PARSEONE_ASSERT(4);
        object->val.map.pair_count = buf_to_u32(&buf[p]);
        p += 4;
        break;

    default:
        if ((object->type >= EMPACK_TYPE_POS_FIXINT) && (object->type <= EMPACK_POS_FIXINT_MAX)) {
            object->val.uint8 = (uint8_t)object->type;
            object->type = EMPACK_TYPE_POS_FIXINT;
        }
        else if ((object->type >= EMPACK_TYPE_FIXMAP) && (object->type <= EMPACK_FIXMAP_MAX)) {
            object->val.map.pair_count = ((uint8_t)object->type) & 0x0f;
            object->type = EMPACK_TYPE_FIXMAP;
        }
        else if ((object->type >= EMPACK_TYPE_FIXSTR) && (object->type <= EMPACK_FIXSTR_MAX)) {
            object->val.str.size = ((uint8_t)object->type) & 0x1f;
            object->type = EMPACK_TYPE_FIXSTR;
        }
        else if ((object->type >= EMPACK_TYPE_NEG_FIXINT) && (object->type <= EMPACK_NEG_FIXINT_MAX)) {
            object->val.int8 = (int8_t)object->type;
            object->type = EMPACK_TYPE_NEG_FIXINT;
        }
        else {
            object->type = EMPACK_TYPE_ERROR;
            ret = EMPACK_STATUS_ERROR;
            goto empack_parse_one_return;
        }
        break;
    }

empack_parse_one_return:
    if ((ret == EMPACK_STATUS_EXTRA) && (p == (*len))) {
        ret = EMPACK_STATUS_COMPLETE;
    }
    (*len) = p;
    return ret;
}

static void *
empack_allocate(empack_t *empack, uint32_t size)
{
    void *ret;

    if (size > (empack->arena_sz - empack->arena_used)) {
        return NULL;
    }

    ret = &((uint8_t *)empack->arena)[empack->arena_used];
    empack->arena_used += size;

    return ret;
}

static void
memswap(uint8_t *dst, const uint8_t *src, uint32_t len)
{
    for (uint32_t i = 0; i < len; i++) {
        dst[i] = src[len - i - 1];
    }
}

static uint16_t
buf_to_u16(const uint8_t *buf)
{
    uint16_t ret;
    memswap((uint8_t *)&ret, buf, 2);
    return ret;
}

static uint32_t
buf_to_u32(const uint8_t *buf)
{
    uint32_t ret;
    memswap((uint8_t *)&ret, buf, 4);
    return ret;
}

static uint64_t
buf_to_u64(const uint8_t *buf)
{
    uint64_t ret;
    memswap((uint8_t *)&ret, buf, 8);
    return ret;
}

static int16_t
buf_to_i16(const uint8_t *buf)
{
    int16_t ret;
    memswap((uint8_t *)&ret, buf, 2);
    return ret;
}

static int32_t
buf_to_i32(const uint8_t *buf)
{
    int32_t ret;
    memswap((uint8_t *)&ret, buf, 4);
    return ret;
}

static int64_t
buf_to_i64(const uint8_t *buf)
{
    int64_t ret;
    memswap((uint8_t *)&ret, buf, 8);
    return ret;
}

static float
buf_to_float(const uint8_t *buf)
{
    float ret;
    memswap((uint8_t *)&ret, buf, 4);
    return ret;
}

static double
buf_to_double(const uint8_t *buf)
{
    double ret;
    memswap((uint8_t *)&ret, buf, 8);
    return ret;
}