#include "config.h"
#include "hpack.h"
#include "hpack-data.h"

struct cno_header_table_t {
    struct cno_header_table_t *prev;
    struct cno_header_table_t *next;
    size_t k_size;
    size_t v_size;
    size_t refcnt;
    char data[];
};

void cno_hpack_free_header(struct cno_header_t *h) {
    if (h->flags & CNO_HEADER_OWNS_NAME)
        free((void *) h->name.data);
    if (h->flags & CNO_HEADER_OWNS_VALUE)
        free((void *) h->value.data);
    if (h->flags & CNO_HEADER_REFS_TABLE) {
        struct cno_header_table_t *entry = ((struct cno_header_table_t *)h->name.data) - 1;
        if (!--entry->refcnt)
            free(entry);
    }
    *h = CNO_HEADER_EMPTY;
}

void cno_hpack_init(struct cno_hpack_t *state, uint32_t limit) {
    state->last = state->first = (struct cno_header_table_t *) state;
    state->size = 0;
    state->limit            = \
    state->limit_upper      = \
    state->limit_update_min = \
    state->limit_update_end = limit;
}

static void cno_hpack_evict(struct cno_hpack_t *state, uint32_t limit) {
    while (state->size > limit) {
        struct cno_header_table_t *entry = state->last;
        state->size -= entry->k_size + entry->v_size + 32;
        entry->next->prev = entry->prev;
        entry->prev->next = entry->next;
        if (!--entry->refcnt)
            free(entry);
    }
}

void cno_hpack_clear(struct cno_hpack_t *state) {
    cno_hpack_evict(state, 0);
}

int cno_hpack_setlimit(struct cno_hpack_t *state, uint32_t limit) {
    if (limit > state->limit_upper)
        return CNO_ERROR(ASSERTION, "dynamic table size limit higher than allowed by peer");
    // The update will be encoded the next time we send a header block.
    if (state->limit_update_min > limit)
        cno_hpack_evict(state, state->limit_update_min = limit);
    state->limit_update_end = limit;
    return CNO_OK;
}

static int cno_hpack_insert(struct cno_hpack_t *state, const struct cno_header_t *h) {
    size_t recorded = h->name.size + h->value.size + 32;
    if (recorded > state->limit) {
        cno_hpack_evict(state, 0);
    } else {
        cno_hpack_evict(state, state->limit - recorded);

        size_t actual = h->name.size + h->value.size + sizeof(struct cno_header_table_t);
        struct cno_header_table_t *entry = malloc(actual);
        if (entry == NULL)
            return CNO_ERROR(NO_MEMORY, "%zu bytes", actual);

        memcpy(&entry->data[0],            h->name.data,  entry->k_size = h->name.size);
        memcpy(&entry->data[h->name.size], h->value.data, entry->v_size = h->value.size);
        entry->refcnt = 1;
        entry->prev = (struct cno_header_table_t *) state;
        entry->next = state->first;
        state->first->prev = entry;
        state->first = entry;
        state->size += recorded;
    }
    return CNO_OK;
}

static int cno_hpack_lookup(struct cno_hpack_t *state, size_t index, struct cno_header_t *out) {
    if (index == 0)
        return CNO_ERROR(PROTOCOL, "header index 0 is reserved");

    if (index <= CNO_HPACK_STATIC_TABLE_SIZE) {
        out->name  = CNO_HPACK_STATIC_TABLE[index - 1].name;
        out->value = CNO_HPACK_STATIC_TABLE[index - 1].value;
        return CNO_OK;
    }

    struct cno_header_table_t *hdr = (struct cno_header_table_t *) state;
    for (index -= CNO_HPACK_STATIC_TABLE_SIZE; index--;)
        if ((hdr = hdr->next) == (struct cno_header_table_t *) state)
            return CNO_ERROR(PROTOCOL, "dynamic table index out of bounds");

    out->name  = (struct cno_buffer_t){ &hdr->data[0], hdr->k_size };
    out->value = (struct cno_buffer_t){ &hdr->data[hdr->k_size], hdr->v_size };
    out->flags |= CNO_HEADER_REFS_TABLE;
    hdr->refcnt++;
    return CNO_OK;
}

// Return value is either 0 (not found), index (name match), or -index (full match).
static int cno_hpack_lookup_inverse(struct cno_hpack_t *state, const struct cno_header_t *needle) {
#define TRY(k, v) do {                   \
    if (!cno_buffer_eq(needle->name, k)) \
        break;                           \
    if (cno_buffer_eq(needle->value, v)) \
        return -i;                       \
    if (possible == 0)                   \
        possible = i;                    \
} while (0)

    int i = 1, possible = 0;
    for (const struct cno_header_t *h = CNO_HPACK_STATIC_TABLE; i <= CNO_HPACK_STATIC_TABLE_SIZE; ++h, ++i)
        TRY(h->name, h->value);
    for (const struct cno_header_table_t *t = state->first; t != (struct cno_header_table_t *) state; t = t->next, ++i)
        TRY(((struct cno_buffer_t) { &t->data[0], t->k_size }),
            ((struct cno_buffer_t) { &t->data[t->k_size], t->v_size }));
    return possible;

#undef TRY
}

// Format: the value as 1 byte if less than mask, else {mask, 0x80 | <7 bits>, ...,
// <last 7 bits>} (little-endian).
static int cno_hpack_decode_uint(struct cno_buffer_t *source, uint8_t mask, size_t *out) {
    if (!source->size)
        return CNO_ERROR(PROTOCOL, "expected uint, got EOF");

    const uint8_t *src = (const uint8_t *) source->data;
    const uint8_t head = *out = *src++ & mask;
    uint8_t size = 1;

    if (head == mask) {
        do {
            if (size == source->size)
                return CNO_ERROR(PROTOCOL, "truncated multi-byte uint");
            if (size == sizeof(size_t))
                return CNO_ERROR(PROTOCOL, "uint literal too large");
            *out += (*src & 0x7Ful) << (7 * size++ - 7);
        } while (*src++ & 0x80);
    }

    *source = cno_buffer_shift(*source, size);
    return CNO_OK;
}

// Format: 1 bit is a flag for Huffman encoding, then a varint for length, then raw data.
// In HPACK, mask is always 0xFF, i.e. a string starts on a byte boundary. QPACK allows
// unaligned strings.
static int cno_hpack_decode_string(struct cno_buffer_t *source, uint8_t mask,
                                   struct cno_buffer_t *out, int *borrow) {
    if (!source->size)
        return CNO_ERROR(PROTOCOL, "expected string, got EOF");
    const uint8_t huffman = (* (const uint8_t *) source->data) & (mask ^ (mask >> 1));

    size_t length = 0;
    if (cno_hpack_decode_uint(source, mask >> 1, &length))
        return CNO_ERROR_UP();
    if (length > source->size)
        return CNO_ERROR(PROTOCOL, "expected %zu octets, got %zu", length, source->size);

    if (length && huffman) {
        uint8_t *buf = malloc(length * 8 / CNO_HUFFMAN_MIN_BITS_PER_CHAR);
        uint8_t *ptr = buf;
        if (!buf)
            return CNO_ERROR(NO_MEMORY, "%zu bytes", length * 2);

        struct cno_huffman_state_t state = CNO_HUFFMAN_STATE_INIT;
        for (const uint8_t *p = (const uint8_t *) source->data, *e = length + p; p != e; p++) {
            state = CNO_HUFFMAN_STATE[state.next << 8 | *p];
            if (state.emit1)
                *ptr++ = state.byte1;
            if (state.emit2)
                *ptr++ = state.byte2;
        }
        if (!state.accept) {
            free(buf);
            return CNO_ERROR(PROTOCOL, "invalid or truncated Huffman code");
        }

        out->data = (char *) buf;
        out->size = ptr - buf;
    } else {
        out->data = source->data;
        out->size = length;
        *borrow = 1;
    }

    *source = cno_buffer_shift(*source, length);
    return CNO_OK;
}

static int cno_hpack_decode_one(struct cno_hpack_t  *state,
                                struct cno_buffer_t *source,
                                struct cno_header_t *target)
{
    *target = CNO_HEADER_EMPTY;

    const uint8_t head = * (const uint8_t *) source->data;
    size_t index = 0;
    if (head >= 0x80) {
        // 1....... -- name & value taken from the table
        return cno_hpack_decode_uint(source, 0x7F, &index)
            || cno_hpack_lookup(state, index, target);
    } else if (head >= 0x40) {
        // 01...... -- name taken from the table, value included as a literal
        if (cno_hpack_decode_uint(source, 0x3F, &index))
            return CNO_ERROR_UP();
    } else if (head >= 0x20) {
        // 001..... -- table size limit update; see cno_hpack_decode
        return CNO_ERROR(PROTOCOL, "unexpected table size limit update");
    } else {
        // 000x.... -- same as 0x40, but we shouldn't insert this header into the table.
        // `x` is set if proxies must use the same encoding.
        target->flags |= CNO_HEADER_NOT_INDEXED;
        if (cno_hpack_decode_uint(source, 0x0F, &index))
            return CNO_ERROR_UP();
    }

    if (index == 0) {
        int borrow = 0;
        if (cno_hpack_decode_string(source, 0xFF, &target->name, &borrow))
            return CNO_ERROR_UP();
        if (!borrow)
            target->flags |= CNO_HEADER_OWNS_NAME;
    } else {
        if (cno_hpack_lookup(state, index, target))
            return CNO_ERROR_UP();
    }

    int borrow = 0;
    if (cno_hpack_decode_string(source, 0xFF, &target->value, &borrow))
        return CNO_ERROR_UP();
    if (!borrow)
        target->flags |= CNO_HEADER_OWNS_VALUE;

    return target->flags & CNO_HEADER_NOT_INDEXED ? CNO_OK : cno_hpack_insert(state, target);
}

int cno_hpack_decode(struct cno_hpack_t *state, struct cno_buffer_t buf, struct cno_header_t *rs, size_t *n) {
    while (buf.size && ((* (const uint8_t *) buf.data) & 0xE0) == 0x20) {
        // 001..... -- a new size limit for the table
        size_t limit = 0;
        if (cno_hpack_decode_uint(&buf, 0x1F, &limit))
            return CNO_ERROR_UP();
        cno_hpack_evict(state, state->limit = limit);
    }

    if (state->limit > state->limit_upper)
        return CNO_ERROR(PROTOCOL, "current decoder state size limit is higher than the upper bound");

    size_t read = 0, limit = *n;
    while (buf.size && read < limit) {
        if (!cno_hpack_decode_one(state, &buf, &rs[read++]))
            continue;
        while (read)
            cno_hpack_free_header(&rs[--read]);
        return CNO_ERROR_UP();
    }
    if (buf.size) {
        while (read)
            cno_hpack_free_header(&rs[--read]);
        return CNO_ERROR(PROTOCOL, "header list too long");
    }
    *n = read;
    return CNO_OK;
}

static int cno_hpack_encode_uint(struct cno_buffer_dyn_t *buf, uint8_t prefix, uint8_t mask, size_t num) {
    if (num < mask) {
        prefix |= num;
        return cno_buffer_dyn_concat(buf, (struct cno_buffer_t) { (char *) &prefix, 1 });
    }

    uint8_t tmp[sizeof(num) * 2], *ptr = tmp;
    *ptr++ = prefix | mask;
    for (num -= mask; num > 0x7F; num >>= 7)
        *ptr++ = num | 0x80;
    *ptr++ = num;
    return cno_buffer_dyn_concat(buf, (struct cno_buffer_t) { (char *) tmp, ptr - tmp });
}

static int cno_hpack_encode_string(struct cno_buffer_dyn_t *buf, uint8_t prefix, uint8_t mask,
                                   const struct cno_buffer_t s) {
    size_t total = 0;
    for (const uint8_t *p = (const uint8_t *) s.data, *e = p + s.size; p != e; p++)
        total += CNO_HUFFMAN_LEN[*p];
    total = (total + 7) / 8;

    if (total >= s.size)
        return cno_hpack_encode_uint(buf, prefix, mask >> 1, s.size) || cno_buffer_dyn_concat(buf, s);

    if (cno_hpack_encode_uint(buf, prefix | (mask ^ (mask >> 1)), mask >> 1, total)
     || cno_buffer_dyn_reserve(buf, buf->size + total))
        return CNO_ERROR_UP();

    uint8_t *out = (uint8_t *) buf->data + buf->size;
    uint64_t code = 0;
    int bits = 0;
    for (const uint8_t *p = (const uint8_t *) s.data, *e = p + s.size; p != e; p++) {
        code  = CNO_HUFFMAN_ENC[*p] | code << CNO_HUFFMAN_LEN[*p];
        bits += CNO_HUFFMAN_LEN[*p];
        if (bits >= 32) {
            uint32_t part = code >> (bits -= 32);
            // FIXME? this is not compiled into a single move because of misalignment
            *out++ = part >> 24;
            *out++ = part >> 16;
            *out++ = part >> 8;
            *out++ = part;
        }
    }
    // At most 31 bits left to flush
    if (bits >= 8) *out++ = code >> (bits -= 8);
    if (bits >= 8) *out++ = code >> (bits -= 8);
    if (bits >= 8) *out++ = code >> (bits -= 8);
    // The remaining byte must be padded with ones
    if (bits) *out++ = 0xff >> bits | code << (8 - bits);

    buf->size += total;
    return CNO_OK;
}

static int cno_hpack_encode_one(struct cno_hpack_t *state, struct cno_buffer_dyn_t *buf, const struct cno_header_t *h) {
    int index = cno_hpack_lookup_inverse(state, h);
    if (index < 0)
        return cno_hpack_encode_uint(buf, 0x80, 0x7F, -index);

    if (h->flags & CNO_HEADER_NOT_INDEXED
        ? cno_hpack_encode_uint(buf, 0x10, 0x0F, index)
        : cno_hpack_encode_uint(buf, 0x40, 0x3F, index) || cno_hpack_insert(state, h))
            return CNO_ERROR_UP();

    if (!index && cno_hpack_encode_string(buf, 0, 0xFF, h->name))
        return CNO_ERROR_UP();

    return cno_hpack_encode_string(buf, 0, 0xFF, h->value);
}

int cno_hpack_encode(struct cno_hpack_t *state, struct cno_buffer_dyn_t *buf,
                     const struct cno_header_t *headers, size_t n)
{
    // Force the other side to evict the same number of entries first...
    if (state->limit != state->limit_update_min)
        if (cno_hpack_encode_uint(buf, 0x20, 0x1F, state->limit = state->limit_update_min))
            return CNO_ERROR_UP();
    // ...then set the limit to its actual value.
    if (state->limit != state->limit_update_end)
        if (cno_hpack_encode_uint(buf, 0x20, 0x1F, state->limit = state->limit_update_min = state->limit_update_end))
            return CNO_ERROR_UP();

    while (n--)
        if (cno_hpack_encode_one(state, buf, headers++))
            return CNO_ERROR_UP();
    return CNO_OK;
}
