#include "w3cs_packed_position.h"

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

static bool selected(const uint8_t *mask, uint32_t index)
{
    return !mask || (mask[index >> 3] & (uint8_t)(1u << (index & 7u)));
}

static uint32_t zigzag_encode(int32_t value)
{
    if (value < 0)
        return (uint32_t)(-(int64_t)value * 2 - 1);
    return (uint32_t)value * 2u;
}

static int32_t zigzag_decode(uint32_t value)
{
    int32_t magnitude = (int32_t)(value >> 1);
    return (value & 1u) ? -magnitude - 1 : magnitude;
}

static uint8_t required_width(uint32_t value)
{
    uint8_t width = 0;
    while (value) {
        ++width;
        value >>= 1;
    }
    return width;
}

static bool append_block(const uint32_t *values, uint32_t count,
                         uint8_t *output, size_t *offset)
{
    uint32_t maximum = 0;
    uint64_t bit_count;
    size_t byte_count;
    size_t block_size;
    uint8_t width;
    uint32_t index;
    if (!count || count > W3CS_PACKED_POSITION_BLOCK_VALUES)
        return false;
    for (index = 0; index < count; ++index)
        if (values[index] > maximum)
            maximum = values[index];
    width = required_width(maximum);
    if (width > W3CS_PACKED_POSITION_MAX_WIDTH)
        return false;
    bit_count = (uint64_t)count * width;
    byte_count = (size_t)((bit_count + 7u) / 8u);
    if (byte_count > SIZE_MAX - 1u || *offset > SIZE_MAX - 1u - byte_count)
        return false;
    block_size = 1u + byte_count;
    if (output) {
        size_t data_offset = *offset + 1u;
        uint64_t bit_offset = 0;
        output[*offset] = width;
        if (byte_count)
            memset(output + data_offset, 0, byte_count);
        for (index = 0; index < count; ++index) {
            uint8_t bit;
            for (bit = 0; bit < width; ++bit) {
                if (values[index] & (1u << bit)) {
                    uint64_t target = bit_offset + bit;
                    output[data_offset + (size_t)(target >> 3)] |=
                        (uint8_t)(1u << (target & 7u));
                }
            }
            bit_offset += width;
        }
    }
    *offset += block_size;
    return true;
}

static bool encode_impl(const int16_t *quanta, uint32_t total_count,
                        uint32_t axis_count, const uint8_t *selection_mask,
                        unsigned drop_bits, uint8_t *output,
                        size_t *output_size)
{
    int32_t predictor[W3CS_PACKED_POSITION_MAX_AXES] = {0, 0, 0, 0};
    uint32_t block[W3CS_PACKED_POSITION_BLOCK_VALUES];
    uint32_t block_count = 0;
    int32_t scale;
    uint32_t index;
    size_t offset = 0;
    if (!output_size || axis_count == 0u
            || axis_count > W3CS_PACKED_POSITION_MAX_AXES
            || drop_bits > 14u || (total_count && !quanta))
        return false;
    scale = (int32_t)(1u << drop_bits);
    for (index = 0; index < total_count; ++index) {
        uint32_t axis = index % axis_count;
        int32_t current;
        int32_t residual;
        if (!selected(selection_mask, index)) {
            predictor[axis] = 0;
            continue;
        }
        if ((int32_t)quanta[index] % scale != 0)
            return false;
        current = (int32_t)quanta[index] / scale;
        residual = current - predictor[axis];
        predictor[axis] = current;
        block[block_count++] = zigzag_encode(residual);
        if (block_count == W3CS_PACKED_POSITION_BLOCK_VALUES) {
            if (!append_block(block, block_count, output, &offset))
                return false;
            block_count = 0;
        }
    }
    if (block_count && !append_block(block, block_count, output, &offset))
        return false;
    *output_size = offset;
    return true;
}

bool w3cs_packed_position_encode(
    const int16_t *quanta,
    uint32_t total_count,
    uint32_t axis_count,
    const uint8_t *selection_mask,
    unsigned drop_bits,
    uint8_t *output,
    size_t output_capacity,
    size_t *output_size)
{
    size_t required = 0;
    if (!encode_impl(quanta, total_count, axis_count, selection_mask,
                     drop_bits, NULL, &required))
        return false;
    if (output_size)
        *output_size = required;
    else
        return false;
    if (!output)
        return true;
    if (output_capacity < required)
        return false;
    return encode_impl(quanta, total_count, axis_count, selection_mask,
                       drop_bits, output, output_size);
}

static bool sparse_padding_is_zero(const uint8_t *selection_mask,
                                   uint32_t total_count)
{
    uint32_t used;
    uint8_t valid;
    if (!selection_mask || !total_count || (total_count & 7u) == 0u)
        return true;
    used = total_count & 7u;
    valid = (uint8_t)((1u << used) - 1u);
    return (selection_mask[total_count >> 3] & (uint8_t)~valid) == 0u;
}

bool w3cs_packed_position_decode(
    const uint8_t *input,
    size_t input_size,
    int16_t *quanta,
    uint32_t total_count,
    uint32_t axis_count,
    const uint8_t *selection_mask,
    unsigned drop_bits,
    size_t *consumed_size)
{
    int32_t predictor[W3CS_PACKED_POSITION_MAX_AXES] = {0, 0, 0, 0};
    uint32_t selected_remaining = 0;
    uint32_t block_count = 0;
    uint32_t block_index = 0;
    uint32_t block_maximum = 0;
    uint8_t block_width = 0;
    size_t block_data_offset = 0;
    int32_t scale;
    int32_t minimum;
    int32_t maximum;
    uint32_t index;
    size_t offset = 0;
    if (!consumed_size || axis_count == 0u
            || axis_count > W3CS_PACKED_POSITION_MAX_AXES
            || drop_bits > 14u || (total_count && !quanta)
            || !sparse_padding_is_zero(selection_mask, total_count))
        return false;
    for (index = 0; index < total_count; ++index)
        if (selected(selection_mask, index))
            ++selected_remaining;
    if (selected_remaining && !input)
        return false;
    scale = (int32_t)(1u << drop_bits);
    minimum = INT16_MIN / scale;
    maximum = INT16_MAX / scale;
    for (index = 0; index < total_count; ++index) {
        uint32_t axis = index % axis_count;
        uint32_t code = 0;
        int32_t residual;
        int32_t current;
        uint8_t bit;
        if (!selected(selection_mask, index)) {
            quanta[index] = 0;
            predictor[axis] = 0;
            continue;
        }
        if (block_index == block_count) {
            uint64_t bits;
            size_t bytes;
            uint8_t padding;
            block_count = selected_remaining <
                    W3CS_PACKED_POSITION_BLOCK_VALUES
                ? selected_remaining : W3CS_PACKED_POSITION_BLOCK_VALUES;
            block_index = 0;
            if (offset >= input_size)
                return false;
            block_width = input[offset++];
            if (block_width > W3CS_PACKED_POSITION_MAX_WIDTH)
                return false;
            bits = (uint64_t)block_count * block_width;
            bytes = (size_t)((bits + 7u) / 8u);
            if (bytes > input_size - offset)
                return false;
            block_data_offset = offset;
            block_maximum = 0;
            offset += bytes;
            padding = (uint8_t)(bits & 7u);
            if (padding && bytes
                    && (input[offset - 1u] &
                        (uint8_t)~((1u << padding) - 1u)) != 0u)
                return false;
        }
        for (bit = 0; bit < block_width; ++bit) {
            uint64_t source = (uint64_t)block_index * block_width + bit;
            if (input[block_data_offset + (size_t)(source >> 3)]
                    & (uint8_t)(1u << (source & 7u)))
                code |= 1u << bit;
        }
        residual = zigzag_decode(code);
        if (code > block_maximum)
            block_maximum = code;
        current = predictor[axis] + residual;
        if (current < minimum || current > maximum)
            return false;
        predictor[axis] = current;
        quanta[index] = (int16_t)(current * scale);
        ++block_index;
        --selected_remaining;
        if (block_index == block_count
                && required_width(block_maximum) != block_width)
            return false;
    }
    *consumed_size = offset;
    return selected_remaining == 0u && block_index == block_count;
}