#include "w3cs_compact_draw.h"

#include <limits.h>
#include <string.h>

static bool append_uleb(uint64_t value, uint8_t *output, size_t capacity,
                        size_t *offset)
{
    do {
        uint8_t byte = (uint8_t)(value & 0x7fu);
        value >>= 7u;
        if (value)
            byte |= 0x80u;
        if (*offset >= capacity)
            return false;
        output[(*offset)++] = byte;
    } while (value);
    return true;
}

static size_t uleb_size(uint64_t value)
{
    size_t size = 1u;
    while (value >= 0x80u) {
        value >>= 7u;
        ++size;
    }
    return size;
}

static bool read_uleb(const uint8_t *input, size_t size, size_t *offset,
                      uint64_t *value)
{
    size_t start = *offset;
    unsigned shift = 0;
    uint64_t result = 0;
    for (;;) {
        uint8_t byte;
        if (*offset >= size || shift >= 64u)
            return false;
        byte = input[(*offset)++];
        if (shift == 63u && (byte & 0x7eu))
            return false;
        result |= (uint64_t)(byte & 0x7fu) << shift;
        if (!(byte & 0x80u))
            break;
        shift += 7u;
    }
    if (*offset - start != uleb_size(result))
        return false;
    *value = result;
    return true;
}

static uint64_t zigzag_encode(int64_t value)
{
    return value < 0
        ? (uint64_t)(-value) * 2u - 1u
        : (uint64_t)value * 2u;
}

static bool zigzag_decode(uint64_t encoded, int64_t *value)
{
    uint64_t magnitude = encoded >> 1u;
    if (magnitude > (uint64_t)INT64_MAX)
        return false;
    *value = (encoded & 1u)
        ? -(int64_t)magnitude - 1 : (int64_t)magnitude;
    return true;
}

static uint32_t find_handle(const struct w3cs_compact_draw_state *state,
                            uint32_t blob_id)
{
    uint32_t index;
    if (!state || !blob_id)
        return 0;
    for (index = 0; index < state->blob_count; ++index)
        if (state->blob_ids[index] == blob_id)
            return index + 1u;
    return 0;
}

static bool append_blob_ref(const struct w3cs_compact_draw_state *state,
                            uint32_t blob_id, uint8_t *output,
                            size_t capacity, size_t *offset)
{
    uint32_t handle = find_handle(state, blob_id);
    if (handle)
        return append_uleb(handle, output, capacity, offset);
    if (!append_uleb(0, output, capacity, offset)
            || capacity - *offset < sizeof(blob_id))
        return false;
    memcpy(output + *offset, &blob_id, sizeof(blob_id));
    *offset += sizeof(blob_id);
    return true;
}

static bool read_blob_ref(const struct w3cs_compact_draw_state *state,
                          const uint8_t *input, size_t size, size_t *offset,
                          uint32_t *blob_id)
{
    uint64_t handle;
    if (!read_uleb(input, size, offset, &handle))
        return false;
    if (handle) {
        if (handle > state->blob_count)
            return false;
        *blob_id = state->blob_ids[(uint32_t)handle - 1u];
        return true;
    }
    if (size - *offset < sizeof(*blob_id))
        return false;
    memcpy(blob_id, input + *offset, sizeof(*blob_id));
    *offset += sizeof(*blob_id);
    return true;
}

void w3cs_compact_draw_reset(struct w3cs_compact_draw_state *state)
{
    if (state)
        memset(state, 0, sizeof(*state));
}

void w3cs_compact_draw_clear_handles(struct w3cs_compact_draw_state *state)
{
    if (state)
        state->blob_count = 0;
}

void w3cs_compact_draw_observe_blob(struct w3cs_compact_draw_state *state,
                                    uint32_t blob_id)
{
    if (!state || !blob_id || find_handle(state, blob_id)
            || state->blob_count >= W3CS_COMPACT_DRAW_MAX_HANDLES)
        return;
    state->blob_ids[state->blob_count++] = blob_id;
}

void w3cs_compact_draw_accept_full(struct w3cs_compact_draw_state *state,
                                   const struct w3cs_indexed_draw *draw)
{
    if (state && draw)
        state->prior = *draw;
}

size_t w3cs_compact_draw_encode(struct w3cs_compact_draw_state *state,
                                const struct w3cs_indexed_draw *draw,
                                uint8_t *output, size_t capacity)
{
    size_t offset = 0;
    bool ok;
    if (!state || !draw || !output)
        return 0;
    ok = append_uleb(draw->primitive_type, output, capacity, &offset)
        && append_uleb(zigzag_encode(
            (int64_t)draw->base_vertex - state->prior.base_vertex),
            output, capacity, &offset)
        && append_uleb(zigzag_encode(
            (int64_t)draw->minimum_vertex - state->prior.minimum_vertex),
            output, capacity, &offset)
        && append_uleb(zigzag_encode(
            (int64_t)draw->vertex_count - state->prior.vertex_count),
            output, capacity, &offset)
        && append_uleb(zigzag_encode(
            (int64_t)draw->start_index - state->prior.start_index),
            output, capacity, &offset)
        && append_uleb(zigzag_encode(
            (int64_t)draw->primitive_count - state->prior.primitive_count),
            output, capacity, &offset)
        && append_blob_ref(state, draw->vertex_blob,
                           output, capacity, &offset)
        && append_blob_ref(state, draw->index_blob,
                           output, capacity, &offset);
    state->prior = *draw;
    return ok ? offset : 0;
}

static bool apply_delta_u32(uint32_t prior, int64_t delta, uint32_t *value)
{
    int64_t result = (int64_t)prior + delta;
    if (result < 0 || result > (int64_t)UINT32_MAX)
        return false;
    *value = (uint32_t)result;
    return true;
}

bool w3cs_compact_draw_decode(struct w3cs_compact_draw_state *state,
                              const uint8_t *input, size_t size,
                              struct w3cs_indexed_draw *draw)
{
    struct w3cs_indexed_draw result;
    uint64_t encoded;
    int64_t delta;
    int64_t base;
    size_t offset = 0;
    if (!state || !input || !draw
            || !read_uleb(input, size, &offset, &encoded)
            || encoded > UINT32_MAX)
        return false;
    result.primitive_type = (uint32_t)encoded;
    if (!read_uleb(input, size, &offset, &encoded)
            || !zigzag_decode(encoded, &delta))
        return false;
    base = (int64_t)state->prior.base_vertex + delta;
    if (base < INT32_MIN || base > INT32_MAX)
        return false;
    result.base_vertex = (int32_t)base;
#define READ_U32_DELTA(field) \
    do { \
        if (!read_uleb(input, size, &offset, &encoded) \
                || !zigzag_decode(encoded, &delta) \
                || !apply_delta_u32(state->prior.field, delta, \
                                    &result.field)) \
            return false; \
    } while (0)
    READ_U32_DELTA(minimum_vertex);
    READ_U32_DELTA(vertex_count);
    READ_U32_DELTA(start_index);
    READ_U32_DELTA(primitive_count);
#undef READ_U32_DELTA
    if (!read_blob_ref(state, input, size, &offset, &result.vertex_blob)
            || !read_blob_ref(state, input, size, &offset,
                              &result.index_blob)
            || offset != size)
        return false;
    state->prior = result;
    *draw = result;
    return true;
}