/*
 * 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 void* input,
	size_t inputSize, void* output, size_t outputSize, size_t& _compressedSize,
	const BCompressionParameters* parameters)
{
#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 = outputSize;
	int zlibError = compress2((Bytef*)output, &bytesUsed, (const Bytef*)input,
		(uLong)inputSize, compressionLevel);
	if (zlibError != Z_OK)
		return _TranslateZlibError(zlibError);
 
	_compressedSize = (size_t)bytesUsed;
	return B_OK;
#else
	return B_NOT_SUPPORTED;
#endif
}
 
 
status_t
BZlibCompressionAlgorithm::DecompressBuffer(const void* input,
	size_t inputSize, void* output, size_t outputSize,
	size_t& _uncompressedSize, const BDecompressionParameters* parameters)
{
	uLongf bytesUsed = outputSize;
	int zlibError = uncompress((Bytef*)output, &bytesUsed, (const Bytef*)input,
		(uLong)inputSize);
	if (zlibError != Z_OK)
		return _TranslateZlibError(zlibError);
 
	_uncompressedSize = (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;
	}
}