/*
 * Copyright 2014, Ingo Weinhold, ingo_weinhold@gmx.de.
 * Distributed under the terms of the MIT License.
 */


#include <ZlibCompressionAlgorithm.h>

#include <errno.h>
#include <string.h>

#include <algorithm>
#include <new>

#include <zlib.h>

#include <DataIO.h>


// build compression support only for userland
#if !defined(_KERNEL_MODE) && !defined(_BOOT_MODE)
#       define B_ZLIB_COMPRESSION_SUPPORT 1
#endif


static const size_t kMinBufferSize              = 1024;
static const size_t kMaxBufferSize              = 1024 * 1024;
static const size_t kDefaultBufferSize  = 4 * 1024;


static size_t
sanitize_buffer_size(size_t size)
{
        if (size < kMinBufferSize)
                return kMinBufferSize;
        return std::min(size, kMaxBufferSize);
}


// #pragma mark - BZlibCompressionParameters


BZlibCompressionParameters::BZlibCompressionParameters(
        int compressionLevel)
        :
        BCompressionParameters(),
        fCompressionLevel(compressionLevel),
        fBufferSize(kDefaultBufferSize),
        fGzipFormat(false)
{
}


BZlibCompressionParameters::~BZlibCompressionParameters()
{
}


int32
BZlibCompressionParameters::CompressionLevel() const
{
        return fCompressionLevel;
}


void
BZlibCompressionParameters::SetCompressionLevel(int32 level)
{
        fCompressionLevel = level;
}


size_t
BZlibCompressionParameters::BufferSize() const
{
        return fBufferSize;
}


void
BZlibCompressionParameters::SetBufferSize(size_t size)
{
        fBufferSize = sanitize_buffer_size(size);
}


bool
BZlibCompressionParameters::IsGzipFormat() const
{
        return fGzipFormat;
}


void
BZlibCompressionParameters::SetGzipFormat(bool gzipFormat)
{
        fGzipFormat = gzipFormat;
}


// #pragma mark - BZlibDecompressionParameters


BZlibDecompressionParameters::BZlibDecompressionParameters()
        :
        BDecompressionParameters(),
        fBufferSize(kDefaultBufferSize)
{
}


BZlibDecompressionParameters::~BZlibDecompressionParameters()
{
}


size_t
BZlibDecompressionParameters::BufferSize() const
{
        return fBufferSize;
}


void
BZlibDecompressionParameters::SetBufferSize(size_t size)
{
        fBufferSize = sanitize_buffer_size(size);
}


// #pragma mark - CompressionStrategy


#ifdef B_ZLIB_COMPRESSION_SUPPORT


struct BZlibCompressionAlgorithm::CompressionStrategy {
        typedef BZlibCompressionParameters Parameters;

        static const bool kNeedsFinalFlush = true;

        static int Init(z_stream& stream,
                const BZlibCompressionParameters* parameters)
        {
                int32 compressionLevel = B_ZLIB_COMPRESSION_DEFAULT;
                bool gzipFormat = false;
                if (parameters != NULL) {
                        compressionLevel = parameters->CompressionLevel();
                        gzipFormat = parameters->IsGzipFormat();
                }

                return deflateInit2(&stream, compressionLevel,
                        Z_DEFLATED,
                        MAX_WBITS + (gzipFormat ? 16 : 0),
                        MAX_MEM_LEVEL,
                        Z_DEFAULT_STRATEGY);
        }

        static void Uninit(z_stream& stream)
        {
                deflateEnd(&stream);
        }

        static int Process(z_stream& stream, bool flush)
        {
                return deflate(&stream, flush ? Z_FINISH : 0);
        }
};


#endif  // B_ZLIB_COMPRESSION_SUPPORT


// #pragma mark - DecompressionStrategy


struct BZlibCompressionAlgorithm::DecompressionStrategy {
        typedef BZlibDecompressionParameters Parameters;

        static const bool kNeedsFinalFlush = false;

        static int Init(z_stream& stream,
                const BZlibDecompressionParameters* /*parameters*/)
        {
                // auto-detect zlib/gzip header
                return inflateInit2(&stream, 32 + MAX_WBITS);
        }

        static void Uninit(z_stream& stream)
        {
                inflateEnd(&stream);
        }

        static int Process(z_stream& stream, bool flush)
        {
                return inflate(&stream, flush ? Z_FINISH : 0);
        }
};


// #pragma mark - Stream


template<typename BaseClass, typename Strategy>
struct BZlibCompressionAlgorithm::Stream : BaseClass {
        Stream(BDataIO* io)
                :
                BaseClass(io),
                fStreamInitialized(false)
        {
        }

        ~Stream()
        {
                if (fStreamInitialized) {
                        if (Strategy::kNeedsFinalFlush)
                                this->Flush();
                        Strategy::Uninit(fStream);
                }
        }

        status_t Init(const typename Strategy::Parameters* parameters)
        {
                status_t error = this->BaseClass::Init(
                        parameters != NULL ? parameters->BufferSize() : kDefaultBufferSize);
                if (error != B_OK)
                        return error;

                memset(&fStream, 0, sizeof(fStream));

                int zlibError = Strategy::Init(fStream, parameters);
                if (zlibError != Z_OK)
                        return _TranslateZlibError(zlibError);

                fStreamInitialized = true;
                return B_OK;
        }

        virtual status_t ProcessData(const void* input, size_t inputSize,
                void* output, size_t outputSize, size_t& bytesConsumed,
                size_t& bytesProduced)
        {
                return _ProcessData(input, inputSize, output, outputSize,
                        bytesConsumed, bytesProduced, false);
        }

        virtual status_t FlushPendingData(void* output, size_t outputSize,
                size_t& bytesProduced)
        {
                size_t bytesConsumed;
                return _ProcessData(NULL, 0, output, outputSize,
                        bytesConsumed, bytesProduced, true);
        }

        template<typename BaseParameters>
        static status_t Create(BDataIO* io, BaseParameters* _parameters,
                BDataIO*& _stream)
        {
                const typename Strategy::Parameters* parameters
#ifdef _BOOT_MODE
                        = static_cast<const typename Strategy::Parameters*>(_parameters);
#else
                        = dynamic_cast<const typename Strategy::Parameters*>(_parameters);
#endif
                Stream* stream = new(std::nothrow) Stream(io);
                if (stream == NULL)
                        return B_NO_MEMORY;

                status_t error = stream->Init(parameters);
                if (error != B_OK) {
                        delete stream;
                        return error;
                }

                _stream = stream;
                return B_OK;
        }

private:
        status_t _ProcessData(const void* input, size_t inputSize,
                void* output, size_t outputSize, size_t& bytesConsumed,
                size_t& bytesProduced, bool flush)
        {
                fStream.next_in = (Bytef*)input;
                fStream.avail_in = inputSize;
                fStream.next_out = (Bytef*)output;
                fStream.avail_out = outputSize;

                int zlibError = Strategy::Process(fStream, flush);
                if (zlibError != Z_OK) {
                        if (zlibError == Z_STREAM_END) {
                                if (fStream.avail_in != 0)
                                        return B_BAD_DATA;
                        } else
                                return _TranslateZlibError(zlibError);
                }

                bytesConsumed = inputSize - (size_t)fStream.avail_in;
                bytesProduced = outputSize - (size_t)fStream.avail_out;
                return B_OK;
        }

private:
        z_stream        fStream;
        bool            fStreamInitialized;
};


// #pragma mark - BZlibCompressionAlgorithm


BZlibCompressionAlgorithm::BZlibCompressionAlgorithm()
        :
        BCompressionAlgorithm()
{
}


BZlibCompressionAlgorithm::~BZlibCompressionAlgorithm()
{
}


status_t
BZlibCompressionAlgorithm::CreateCompressingInputStream(BDataIO* input,
        const BCompressionParameters* parameters, BDataIO*& _stream)
{
#ifdef B_ZLIB_COMPRESSION_SUPPORT
        return Stream<BAbstractInputStream, CompressionStrategy>::Create(
                input, parameters, _stream);
#else
        return B_NOT_SUPPORTED;
#endif
}


status_t
BZlibCompressionAlgorithm::CreateCompressingOutputStream(BDataIO* output,
        const BCompressionParameters* parameters, BDataIO*& _stream)
{
#ifdef B_ZLIB_COMPRESSION_SUPPORT
        return Stream<BAbstractOutputStream, CompressionStrategy>::Create(
                output, parameters, _stream);
#else
        return B_NOT_SUPPORTED;
#endif
}


status_t
BZlibCompressionAlgorithm::CreateDecompressingInputStream(BDataIO* input,
        const BDecompressionParameters* parameters, BDataIO*& _stream)
{
        return Stream<BAbstractInputStream, DecompressionStrategy>::Create(
                input, parameters, _stream);
}


status_t
BZlibCompressionAlgorithm::CreateDecompressingOutputStream(BDataIO* output,
        const BDecompressionParameters* parameters, BDataIO*& _stream)
{
        return Stream<BAbstractOutputStream, DecompressionStrategy>::Create(
                output, parameters, _stream);
}


status_t
BZlibCompressionAlgorithm::CompressBuffer(const iovec& input, iovec& output,
        const BCompressionParameters* parameters, iovec* scratch)
{
#ifdef B_ZLIB_COMPRESSION_SUPPORT
        const BZlibCompressionParameters* zlibParameters
                = dynamic_cast<const BZlibCompressionParameters*>(parameters);
        int compressionLevel = zlibParameters != NULL
                ? zlibParameters->CompressionLevel()
                : B_ZLIB_COMPRESSION_DEFAULT;

        uLongf bytesUsed = output.iov_len;
        int zlibError = compress2((Bytef*)output.iov_base, &bytesUsed,
                (const Bytef*)input.iov_base, (uLong)input.iov_len, compressionLevel);
        if (zlibError != Z_OK)
                return _TranslateZlibError(zlibError);

        output.iov_len = (size_t)bytesUsed;
        return B_OK;
#else
        return B_NOT_SUPPORTED;
#endif
}


status_t
BZlibCompressionAlgorithm::DecompressBuffer(const iovec& input, iovec& output,
        const BDecompressionParameters* parameters, iovec* scratch)
{
        uLongf bytesUsed = output.iov_len;
        int zlibError = uncompress((Bytef*)output.iov_base, &bytesUsed,
                (const Bytef*)input.iov_base, (uLong)input.iov_len);
        if (zlibError != Z_OK)
                return _TranslateZlibError(zlibError);

        output.iov_len = (size_t)bytesUsed;
        return B_OK;
}


/*static*/ status_t
BZlibCompressionAlgorithm::_TranslateZlibError(int error)
{
        switch (error) {
                case Z_OK:
                        return B_OK;
                case Z_STREAM_END:
                case Z_NEED_DICT:
                        // a special event (no error), but the caller doesn't seem to handle
                        // it
                        return B_ERROR;
                case Z_ERRNO:
                        return errno;
                case Z_STREAM_ERROR:
                        return B_BAD_VALUE;
                case Z_DATA_ERROR:
                        return B_BAD_DATA;
                case Z_MEM_ERROR:
                        return B_NO_MEMORY;
                case Z_BUF_ERROR:
                        return B_BUFFER_OVERFLOW;
                case Z_VERSION_ERROR:
                        return B_BAD_VALUE;
                default:
                        return B_ERROR;
        }
}