#include <zlib.h>
#include <http_common.h>
#include <http_half.h>
#include <http_convert.h>
#include <event2/buffer.h>
#include <brotli/decode.h>
#include <brotli/encode.h>

struct hf_content_uncompress
{
    /* MODE AND CALLBACKS */
    unsigned int content_encode;
    hf_private_cb * data_cb;
    void * data_cb_user;

    /* ZLIB STREAM */
    z_stream * z_stream_ptr;
    /* BR DECODER STATE */
    BrotliDecoderState * brdec_state;

    unsigned char * chunk;
    size_t sz_chunk;
};

struct hf_content_compress
{
    unsigned int content_encode;
    z_stream * z_stream_ptr;
    BrotliEncoderState * brenc_state;
};

void hf_content_uncompress_destroy(struct hf_content_uncompress * cv_object)
{
    if (cv_object->z_stream_ptr != NULL)
    {
        inflateEnd(cv_object->z_stream_ptr);
        free(cv_object->z_stream_ptr);
        cv_object->z_stream_ptr = NULL;
    }

    if (cv_object->brdec_state)
    {
        BrotliDecoderDestroyInstance(cv_object->brdec_state);
        cv_object->brdec_state = NULL;
    }

    free(cv_object->chunk);
    free(cv_object);
}

struct hf_content_uncompress * hf_content_uncompress_create(unsigned int content_encode,
    hf_private_cb * data_cb, void * data_cb_user)
{
    struct hf_content_uncompress * cv_object = ALLOC(struct hf_content_uncompress, 1);
    assert(data_cb != NULL);

    cv_object->content_encode = content_encode;
    cv_object->data_cb = data_cb;
    cv_object->data_cb_user = data_cb_user;

    /* CHUNK, 4K */
#define CHUNK_SIZE    (1024 * 1024 * 4)
    cv_object->chunk = (unsigned char *) malloc(CHUNK_SIZE);
    cv_object->sz_chunk = CHUNK_SIZE;

    /* ZSTREAM */
    if (content_encode == HTTP_ACCEPT_ENCODING_GZIP || content_encode == HTTP_ACCEPT_ENCODING_DEFLATE)
    {
        cv_object->z_stream_ptr = ALLOC(z_stream, 1);
        cv_object->z_stream_ptr->zalloc = NULL;
        cv_object->z_stream_ptr->zfree = NULL;
        cv_object->z_stream_ptr->opaque = NULL;
        cv_object->z_stream_ptr->avail_in = 0;
        cv_object->z_stream_ptr->next_in = Z_NULL;

        int ret = 0;
        if (content_encode == HTTP_ACCEPT_ENCODING_GZIP)
        {
            ret = inflateInit2(cv_object->z_stream_ptr, MAX_WBITS + 16);
        }
        if (content_encode == HTTP_ACCEPT_ENCODING_DEFLATE)
        {
            ret = inflateInit2(cv_object->z_stream_ptr, -MAX_WBITS);
        }
        if (ret != Z_OK)
        {
            goto __errout;
        }
    }

    if (content_encode == HTTP_ACCEPT_ENCODING_BR)
    {
        cv_object->brdec_state = BrotliDecoderCreateInstance(NULL, NULL, NULL);
        if (unlikely(cv_object->brdec_state == NULL)) goto __errout;
    }

    return cv_object;

__errout:
    if (cv_object->z_stream_ptr != NULL)
    {
        free(cv_object->z_stream_ptr);
        cv_object->z_stream_ptr = NULL;
    }

    if (cv_object->brdec_state != NULL)
    {
        BrotliDecoderDestroyInstance(cv_object->brdec_state);
        cv_object->brdec_state = NULL;
    }

    free(cv_object->chunk);
    free(cv_object);
    return NULL;
}

static int __hf_content_uncompress_write_zlib(struct hf_content_uncompress * cv_object,
    struct http_half_private * hf_private, tfe_http_event http_ev, const unsigned char * data, size_t datalen)
{
    z_stream * z_stream_ptr = cv_object->z_stream_ptr;
    z_stream_ptr->avail_in = (unsigned int) datalen;
    z_stream_ptr->next_in = (unsigned char *) data;

    if (z_stream_ptr->avail_in == 0)
    {
        (void) inflateEnd(z_stream_ptr);
        return Z_ERRNO;
    }

    int ret = 0;
    do
    {
        z_stream_ptr->avail_out = (unsigned int) cv_object->sz_chunk;
        z_stream_ptr->next_out = cv_object->chunk;

        ret = inflate(z_stream_ptr, Z_NO_FLUSH);

        if (ret == Z_STREAM_ERROR || ret == Z_NEED_DICT ||
            ret == Z_DATA_ERROR || ret == Z_MEM_ERROR)
        {
            goto __error;
        }

        unsigned int have = (unsigned int) cv_object->sz_chunk - z_stream_ptr->avail_out;
        if (have > 0 && cv_object->data_cb != NULL)
        {
            cv_object->data_cb(hf_private, http_ev, cv_object->chunk, (size_t) have, cv_object->data_cb_user);
        }

    }
    while (z_stream_ptr->avail_out == 0);
    return ret;

__error:
    (void) inflateEnd(z_stream_ptr);
    return ret;
}

static int __hf_content_uncompress_write_br(struct hf_content_uncompress * cv_object,
    struct http_half_private * hf_private, tfe_http_event http_ev, const unsigned char * data, size_t datalen)
{
    size_t available_in = datalen;
    const unsigned char * next_in = data;
    size_t available_out;
    unsigned char * next_out;

    int ret;
    for (;;)
    {
        available_out = cv_object->sz_chunk;
        next_out = cv_object->chunk;

        ret = BrotliDecoderDecompressStream(cv_object->brdec_state, &available_in, &next_in,
            &available_out, &next_out, 0);

        size_t have = (unsigned int) cv_object->sz_chunk - available_out;
        if (have > 0 && cv_object->data_cb != NULL)
        {
            cv_object->data_cb(hf_private, http_ev, cv_object->chunk, have, cv_object->data_cb_user);
        }

        if (ret == BROTLI_DECODER_RESULT_SUCCESS || ret == BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT)
        {
            return 1;
        }

        if (ret == BROTLI_DECODER_RESULT_ERROR)
        {
            BrotliDecoderErrorCode errcode = BrotliDecoderGetErrorCode(cv_object->brdec_state);
            TFE_LOG_ERROR(g_http_plugin->logger, "BrotliDecoderDecompressStream() failed: errno = %d, %s",
                errcode, BrotliDecoderErrorString(errcode)); return -1;
        }

        assert(ret == BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT);
    }
}

int hf_content_uncompress_write(struct hf_content_uncompress * cv_object,
    struct http_half_private * hf_private, tfe_http_event http_ev, const unsigned char * data, size_t datalen)
{
    if (cv_object->content_encode == HTTP_ACCEPT_ENCODING_GZIP ||
        cv_object->content_encode == HTTP_ACCEPT_ENCODING_DEFLATE)
    {
        return __hf_content_uncompress_write_zlib(cv_object, hf_private, http_ev, data, datalen);
    }

    if (cv_object->content_encode == HTTP_ACCEPT_ENCODING_BR)
    {
        return __hf_content_uncompress_write_br(cv_object, hf_private, http_ev, data, datalen);
    }

    assert(0);
    return -1;
}

struct hf_content_compress * hf_content_compress_create(unsigned int content_encode)
{
    struct hf_content_compress * cv_object = ALLOC(struct hf_content_compress, 1);
    cv_object->content_encode = content_encode;

    if (cv_object->content_encode == HTTP_ACCEPT_ENCODING_GZIP ||
        cv_object->content_encode == HTTP_ACCEPT_ENCODING_DEFLATE)
    {
        cv_object->z_stream_ptr = ALLOC(z_stream, 1);
        cv_object->z_stream_ptr->zalloc = NULL;
        cv_object->z_stream_ptr->zfree = NULL;
        cv_object->z_stream_ptr->opaque = NULL;
        cv_object->z_stream_ptr->avail_in = 0;
        cv_object->z_stream_ptr->next_in = Z_NULL;

        int __windows_bits = 0;
        if (content_encode == HTTP_ACCEPT_ENCODING_GZIP)
        {
            __windows_bits = MAX_WBITS + 16;
        }

        if (content_encode == HTTP_ACCEPT_ENCODING_DEFLATE)
        {
            __windows_bits = -MAX_WBITS;
        }

        int ret = deflateInit2(cv_object->z_stream_ptr, Z_DEFAULT_COMPRESSION,
            Z_DEFLATED, __windows_bits, 8, Z_DEFAULT_STRATEGY);

        if (ret != Z_OK) goto __errout;
    }

    if (cv_object->content_encode == HTTP_ACCEPT_ENCODING_BR)
    {
        cv_object->brenc_state = BrotliEncoderCreateInstance(NULL, NULL, NULL);
        if (unlikely(cv_object->brenc_state == NULL)) goto __errout;
    }

    return cv_object;

__errout:
    if (cv_object->z_stream_ptr)
    {
        free(&cv_object->z_stream_ptr);
        cv_object->z_stream_ptr = NULL;
    }

    if (cv_object->brenc_state)
    {
        BrotliEncoderDestroyInstance(cv_object->brenc_state);
        cv_object->brenc_state = NULL;
    }

    free(cv_object);
    return NULL;
}

#define SZ_RESERVE_SPACE                8192

static int __hf_content_compress_write_zlib(struct hf_content_compress * cv_object,
    const unsigned char * in_data, size_t sz_in_data, struct evbuffer * out_ev_buf, int end)
{

    struct evbuffer_iovec v[1];

    /* Reserve the space, because the length of the compressed data will be short
     * than uncompressed data in usually, we set the reserve space as much as sz_in_data */
    size_t __sz_reserve_chunk = sz_in_data > SZ_RESERVE_SPACE ? sz_in_data : SZ_RESERVE_SPACE;
    int iov_count = evbuffer_reserve_space(out_ev_buf, __sz_reserve_chunk, v, 1);
    if (iov_count != 1) return -1;

    z_stream * z = cv_object->z_stream_ptr;
    z->next_in = (unsigned char *) in_data;
    z->avail_in = (unsigned int) sz_in_data;

    z->next_out = (unsigned char *) v[0].iov_base;
    z->avail_out = (unsigned int) v[0].iov_len;

    int flush = end ? Z_FINISH : Z_NO_FLUSH;
    int deflate_ret = 0;
    do
    {
        deflate_ret = deflate(z, flush);
        assert(deflate_ret != Z_STREAM_ERROR);

        if (z->avail_in == 0 || z->avail_out == 0)
        {
            v[0].iov_len = (unsigned int) v[0].iov_len - z->avail_out;
            int ret = evbuffer_commit_space(out_ev_buf, v, iov_count);
            if(ret < 0) return -2;

            /* Need more space */
            if(z->avail_out == 0)
            {
                iov_count = evbuffer_reserve_space(out_ev_buf, __sz_reserve_chunk, v, 1);
                if(unlikely(iov_count != 1)) return -3;

                z->next_out = (unsigned char *) v[0].iov_base;
                z->avail_out = (unsigned int) v[0].iov_len;
            }
        }
    }
    while (z->avail_in > 0);
    assert(end == 0 || deflate_ret == Z_STREAM_END);

    (void) deflate_ret;
    return 0;
}

static int __hf_content_compress_write_br(struct hf_content_compress * cv_object,
    const unsigned char * in_data, size_t sz_in_data, struct evbuffer * out_ev_buf, int end)
{
    struct evbuffer_iovec v[1];

    /* Reserve the space, because the length of the compressed data will be short
     * than uncompressed data in usually, we set the reserve space as much as sz_in_data */
    size_t __sz_reserve_chunk = sz_in_data > SZ_RESERVE_SPACE ? sz_in_data : SZ_RESERVE_SPACE;
    int iov_count = evbuffer_reserve_space(out_ev_buf, __sz_reserve_chunk, v, 1);
    if (iov_count != 1) return -1;

    const unsigned char * next_in = in_data;
    size_t avail_in = sz_in_data;

    unsigned char * next_out = (unsigned char *)v[0].iov_base;
    size_t avail_out = v[0].iov_len;

    enum BrotliEncoderOperation op = end ? BROTLI_OPERATION_FINISH : BROTLI_OPERATION_PROCESS;
    int ret = 0;
    do
    {
        ret = BrotliEncoderCompressStream(cv_object->brenc_state, op,
            &avail_in, &next_in, &avail_out, &next_out, NULL);

        if(unlikely(ret == BROTLI_FALSE))
        {
            TFE_LOG_ERROR(g_http_plugin->logger, "BrotliEncoderCompressStream() error.");
            return ret;
        }

        if (avail_out == 0 || avail_in == 0)
        {
            v[0].iov_len = v[0].iov_len - avail_out;
            ret = evbuffer_commit_space(out_ev_buf, v, iov_count);
            if(ret < 0) return -2;

            if(avail_out == 0)
            {
                iov_count = evbuffer_reserve_space(out_ev_buf, __sz_reserve_chunk, v, 1);
                if(unlikely(iov_count != 1)) return -3;

                next_out = (unsigned char *) v[0].iov_base;
                avail_out = (unsigned int) v[0].iov_len;
            }
        }
    }
    while (avail_in > 0);
    return 0;
}

int hf_content_compress_write(struct hf_content_compress * cv_object,
    const unsigned char * in_data, size_t sz_in_data, struct evbuffer * out_ev_buf, int end)
{
    if (cv_object->content_encode == HTTP_ACCEPT_ENCODING_GZIP ||
        cv_object->content_encode == HTTP_ACCEPT_ENCODING_DEFLATE)
    {
        return __hf_content_compress_write_zlib(cv_object, in_data, sz_in_data, out_ev_buf, end);
    }

    if (cv_object->content_encode == HTTP_ACCEPT_ENCODING_BR)
    {
        return __hf_content_compress_write_br(cv_object, in_data, sz_in_data, out_ev_buf, end);
    }

    assert(0);
    return -1;
}

void hf_content_compress_destroy(hf_content_compress * cv_object)
{
    (void) deflateEnd(cv_object->z_stream_ptr);
    free(cv_object->z_stream_ptr);
    cv_object->z_stream_ptr = NULL;
    free(cv_object);
}