#include "tile/base/compression/gzip.h"

#include "tile/base/compression/util.h"

#include "zlib/zlib.h"

namespace tile {
namespace compression {
TILE_COMPRESSION_REGISTER_COMPRESSOR("gzip", GzipCompressor);
TILE_COMPRESSION_REGISTER_DECOMPRESSOR("gzip", GzipDecompressor);

namespace detail {
struct ZStream {
    z_stream stream;
};
}// namespace detail

namespace {

inline double
EstimateCompressionRate(const z_stream *stream, double default_value)
{
    if (stream->total_in > 0) {
        double rate               = 1.0f * stream->total_out / stream->total_in;
        constexpr double kMinRate = 1.1;
        return rate * kMinRate;
    }

    return default_value;
}

inline uint32_t
RestrictAvailSize(size_t size)
{
    return static_cast<uint32_t>(std::min(size, static_cast<size_t>(UINT32_MAX)));
}

// See document of inflateInit2 in zlib.h
constexpr int ZLIB_INIT_FLAG_GZIP = 16;

// he size to increase every time Z_BUF_ERROR returns.
constexpr int kOutBufferIncreaseSize = 32;

bool
DoAppend(z_stream *stream,
         CompressionOutputStream *out_stream,
         const void *input_buffer,
         std::size_t input_size,
         bool is_deflate,
         bool finish)
{
    TILE_CHECK(stream && out_stream);

    if (!finish && (input_buffer == nullptr || input_size == 0)) {
        // no input data
        return true;
    }

    stream->next_in = reinterpret_cast<Bytef *>(const_cast<void *>(input_buffer));
    int code        = Z_OK;
    // the number of  `buffer too small error`
    int need_more_space_cnt = 0;
    std::string tmp_buffer;
    size_t left_size = input_size;
    while (left_size > 0 || need_more_space_cnt > 0 || finish) {
        stream->avail_in            = RestrictAvailSize(left_size);
        const auto current_avail_in = stream->avail_in;
        std::size_t out_size;
        if (!need_more_space_cnt) {
            void *out_data;
            if (!out_stream->Next(&out_data, &out_size)) {
                // out buffer is full
                return false;
            }

            stream->next_out = reinterpret_cast<Bytef *>(out_data);
        } else {
            double rate = EstimateCompressionRate(stream, 0.5);
            tmp_buffer.resize(left_size * rate + need_more_space_cnt * kOutBufferIncreaseSize);
            stream->next_out = reinterpret_cast<Bytef *>(internal::RemoveConstPtr(tmp_buffer.data()));
            out_size         = tmp_buffer.size();
        }

        stream->avail_out                   = RestrictAvailSize(out_size);
        const std::size_t current_avail_out = stream->avail_out;

        TILE_CHECK_GT(stream->avail_out, 0, "Avail_out should never be zero before the call.");
        int flush_option = finish ? Z_FINISH : Z_NO_FLUSH;
        if (is_deflate) {
            code = deflate(stream, flush_option);
        } else {
            code = inflate(stream, flush_option);
        }

        // https://refspecs.linuxbase.org/LSB_3.0.0/LSB-Core-generic/LSB-Core-generic/zlib-inflate-1.html
        // No progress is possible; either avail_in or avail_out was zero.
        if (code == Z_BUF_ERROR) {
            if (need_more_space_cnt == 0) { out_stream->BackUp(out_size); }

            if (stream->avail_in == 0) {
                // if not finish, we need more input data
                // otherwise, fail to compress/decompress
                return !finish;
            }

            ++need_more_space_cnt;
            continue;
        }

        if (code < 0) {
            TILE_LOG_ERROR_EVERY_SECOND("gzip compress error code [{}]", code);
            break;
        }

        // consume input data (current_avail_in - stream->avail_in)
        // produce output data (stream->avail_out)
        left_size -= current_avail_in - stream->avail_in;
        if (need_more_space_cnt == 0) {
            out_stream->BackUp(stream->avail_out);
        } else {
            if (TILE_UNLIKELY(!CopyDataToCompressionOutputStream(
                    out_stream, tmp_buffer.data(), current_avail_out - stream->avail_out))) {
                return false;
            }
            need_more_space_cnt = 0;
        }

        if (code == Z_STREAM_END) { return true; }
    }

    return code == Z_OK;
}
}// namespace

bool
GzipCompressor::Compress(const void *src, std::size_t size, CompressionOutputStream *out)
{
    bool ok = Init(out);
    ok &= Append(src, size);
    ok &= Flush();
    out_ = nullptr;
    return ok;
}

bool
GzipCompressor::Compress(const NoncontiguousBuffer &bytes, CompressionOutputStream *out)
{
    bool ok          = Init(out);
    std::size_t left = bytes.ByteSize();
    for (auto iter = bytes.begin(); ok && iter != bytes.end(); ++iter) {
        auto len = std::min(left, iter->size());
        ok &= Append(iter->data(), len);
        left -= len;
    }
    ok &= Flush();
    out_ = nullptr;
    return ok;
}

bool
GzipCompressor::Append(const void *buffer, std::size_t size)
{
    return DoAppend(&stream_->stream, out_, buffer, size, true, false);
}

bool
GzipCompressor::Flush()
{
    TILE_CHECK(out_);
    if (!DoAppend(&stream_->stream, out_, nullptr, 0, true, true)) { return false; }
    return Release();
}

bool
GzipCompressor::Init(CompressionOutputStream *out)
{
    TILE_CHECK(!out_);
    stream_  = make_unique<detail::ZStream>();
    int code = deflateInit2(
        &stream_->stream, Z_DEFAULT_COMPRESSION, Z_DEFLATED, MAX_WBITS + ZLIB_INIT_FLAG_GZIP, 8, Z_DEFAULT_STRATEGY);
    if (code != Z_OK) {
        TILE_LOG_ERROR_EVERY_SECOND("deflateInit2 error code [{}]", code);
        stream_ = nullptr;
    } else {
        out_ = out;
    }
    return code == Z_OK;
}

bool
GzipCompressor::Release()
{
    TILE_CHECK(stream_);
    int code = deflateEnd(&stream_->stream);
    if (code != Z_OK) { TILE_LOG_WARNING_EVERY_SECOND("DeflateEnd fail with code [{}].", code); }
    return code == Z_OK;
}

bool
GzipDecompressor::Decompress(const void *src, std::size_t size, CompressionOutputStream *out)
{
    bool ok = Init(out);
    ok &= Append(src, size);
    ok &= Flush();
    out_ = nullptr;
    return ok;
}

bool
GzipDecompressor::Decompress(const NoncontiguousBuffer &compressed, CompressionOutputStream *out)
{
    bool ok          = Init(out);
    std::size_t left = compressed.ByteSize();
    for (auto iter = compressed.begin(); ok && iter != compressed.end(); ++iter) {
        auto len = std::min(left, iter->size());
        ok &= Append(iter->data(), len);
        left -= len;
    }
    ok &= Flush();
    out_ = nullptr;
    return ok;
}

bool
GzipDecompressor::Append(const void *buffer, std::size_t size)
{
    return DoAppend(&stream_->stream, out_, buffer, size, false, false);
}

bool
GzipDecompressor::Flush()
{
    TILE_CHECK(out_);
    if (!DoAppend(&stream_->stream, out_, nullptr, 0, false, true)) { return false; }
    return Release();
}

bool
GzipDecompressor::Init(CompressionOutputStream *out)
{
    TILE_CHECK(!out_);
    stream_  = make_unique<detail::ZStream>();
    int code = inflateInit2(&stream_->stream, MAX_WBITS + ZLIB_INIT_FLAG_GZIP);
    if (code != Z_OK) {
        TILE_LOG_ERROR_EVERY_SECOND("inflateInit2 error code [{}]", code);
        stream_ = nullptr;
    } else {
        out_ = out;
    }
    return code == Z_OK;
}

bool
GzipDecompressor::Release()
{
    TILE_CHECK(stream_);
    int code = inflateEnd(&stream_->stream);
    if (code != Z_OK) { TILE_LOG_WARNING_EVERY_SECOND("DeflateEnd fail with code [{}].", code); }
    return code == Z_OK;
}

}// namespace compression
}// namespace tile
