/**

* Copyright (c) 2016-present, Gregory Szorc

* All rights reserved.

*

* This software may be modified and distributed under the terms

* of the BSD license. See the LICENSE file for details.

*/

#include "python-zstandard.h"

#include "pool.h"

extern PyObject* ZstdError;

int ensure_cctx(ZstdCompressor* compressor) {

	size_t zresult;

	assert(compressor);

	assert(compressor->cctx);

	assert(compressor->params);

	ZSTD_CCtx_reset(compressor->cctx);

	zresult = ZSTD_CCtx_setParametersUsingCCtxParams(compressor->cctx, compressor->params);

	if (ZSTD_isError(zresult)) {

		PyErr_Format(ZstdError, "could not set compression parameters: %s",

			ZSTD_getErrorName(zresult));

		return 1;

	}

	if (compressor->dict) {

		if (compressor->dict->cdict) {

			zresult = ZSTD_CCtx_refCDict(compressor->cctx, compressor->dict->cdict);

		}

		else {

			zresult = ZSTD_CCtx_loadDictionary_advanced(compressor->cctx,

				compressor->dict->dictData, compressor->dict->dictSize,

				ZSTD_dlm_byRef, compressor->dict->dictType);

		}

		if (ZSTD_isError(zresult)) {

			PyErr_Format(ZstdError, "could not load compression dictionary: %s",

				ZSTD_getErrorName(zresult));

			return 1;

		}

	}

	return 0;

}

static PyObject* frame_progression(ZSTD_CCtx* cctx) {

	PyObject* result = NULL;

	PyObject* value;

	ZSTD_frameProgression progression;

	result = PyTuple_New(3);

	if (!result) {

		return NULL;

	}

	progression = ZSTD_getFrameProgression(cctx);

	value = PyLong_FromUnsignedLongLong(progression.ingested);

	if (!value) {

		Py_DECREF(result);

		return NULL;

	}

	PyTuple_SET_ITEM(result, 0, value);

	value = PyLong_FromUnsignedLongLong(progression.consumed);

	if (!value) {

		Py_DECREF(result);

		return NULL;

	}

	PyTuple_SET_ITEM(result, 1, value);

	value = PyLong_FromUnsignedLongLong(progression.produced);

	if (!value) {

		Py_DECREF(result);

		return NULL;

	}

	PyTuple_SET_ITEM(result, 2, value);

	return result;

}

PyDoc_STRVAR(ZstdCompressor__doc__,

"ZstdCompressor(level=None, dict_data=None, compression_params=None)\n"

"\n"

"Create an object used to perform Zstandard compression.\n"

"\n"

"An instance can compress data various ways. Instances can be used multiple\n"

"times. Each compression operation will use the compression parameters\n"

"defined at construction time.\n"

"\n"

"Compression can be configured via the following names arguments:\n"

"\n"

"level\n"

"   Integer compression level.\n"

"dict_data\n"

"   A ``ZstdCompressionDict`` to be used to compress with dictionary data.\n"

"compression_params\n"

"   A ``CompressionParameters`` instance defining low-level compression"

"   parameters. If defined, this will overwrite the ``level`` argument.\n"

"write_checksum\n"

"   If True, a 4 byte content checksum will be written with the compressed\n"

"   data, allowing the decompressor to perform content verification.\n"

"write_content_size\n"

"   If True (the default), the decompressed content size will be included in\n"

"   the header of the compressed data. This data will only be written if the\n"

"   compressor knows the size of the input data.\n"

"write_dict_id\n"

"   Determines whether the dictionary ID will be written into the compressed\n"

"   data. Defaults to True. Only adds content to the compressed data if\n"

"   a dictionary is being used.\n"

"threads\n"

"   Number of threads to use to compress data concurrently. When set,\n"

"   compression operations are performed on multiple threads. The default\n"

"   value (0) disables multi-threaded compression. A value of ``-1`` means to\n"

"   set the number of threads to the number of detected logical CPUs.\n"

);

static int ZstdCompressor_init(ZstdCompressor* self, PyObject* args, PyObject* kwargs) {

	static char* kwlist[] = {

		"level",

		"dict_data",

		"compression_params",

		"write_checksum",

		"write_content_size",

		"write_dict_id",

		"threads",

		NULL

	};

	int level = 3;

	ZstdCompressionDict* dict = NULL;

	ZstdCompressionParametersObject* params = NULL;

	PyObject* writeChecksum = NULL;

	PyObject* writeContentSize = NULL;

	PyObject* writeDictID = NULL;

	int threads = 0;

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|iO!O!OOOi:ZstdCompressor",

		kwlist,	&level, &ZstdCompressionDictType, &dict,

		&ZstdCompressionParametersType, &params,

		&writeChecksum, &writeContentSize, &writeDictID, &threads)) {

		return -1;

	}

	if (level > ZSTD_maxCLevel()) {

		PyErr_Format(PyExc_ValueError, "level must be less than %d",

			ZSTD_maxCLevel() + 1);

		return -1;

	}

	if (threads < 0) {

		threads = cpu_count();

	}

	/* We create a ZSTD_CCtx for reuse among multiple operations to reduce the

	   overhead of each compression operation. */

	self->cctx = ZSTD_createCCtx();

	if (!self->cctx) {

		PyErr_NoMemory();

		return -1;

	}

	/* TODO stuff the original parameters away somewhere so we can reset later. This

	   will allow us to do things like automatically adjust cparams based on input

	   size (assuming zstd isn't doing that internally). */

	self->params = ZSTD_createCCtxParams();

	if (!self->params) {

		PyErr_NoMemory();

		return -1;

	}

	if (params && writeChecksum) {

		PyErr_SetString(PyExc_ValueError,

			"cannot define compression_params and write_checksum");

		return -1;

	}

	if (params && writeContentSize) {

		PyErr_SetString(PyExc_ValueError,

			"cannot define compression_params and write_content_size");

		return -1;

	}

	if (params && writeDictID) {

		PyErr_SetString(PyExc_ValueError,

			"cannot define compression_params and write_dict_id");

		return -1;

	}

	if (params && threads) {

		PyErr_SetString(PyExc_ValueError,

			"cannot define compression_params and threads");

		return -1;

	}

	if (params) {

		if (set_parameters(self->params, params)) {

			return -1;

		}

	}

	else {

		if (set_parameter(self->params, ZSTD_p_compressionLevel, level)) {

			return -1;

		}

		if (set_parameter(self->params, ZSTD_p_contentSizeFlag,

			writeContentSize ? PyObject_IsTrue(writeContentSize) : 1)) {

			return -1;

		}

		if (set_parameter(self->params, ZSTD_p_checksumFlag,

			writeChecksum ? PyObject_IsTrue(writeChecksum) : 0)) {

			return -1;

		}

		if (set_parameter(self->params, ZSTD_p_dictIDFlag,

			writeDictID ? PyObject_IsTrue(writeDictID) : 1)) {

			return -1;

		}

		if (threads) {

			if (set_parameter(self->params, ZSTD_p_nbWorkers, threads)) {

				return -1;

			}

		}

	}

	if (dict) {

		self->dict = dict;

		Py_INCREF(dict);

	}

	if (ensure_cctx(self)) {

		return -1;

	}

	return 0;

}

static void ZstdCompressor_dealloc(ZstdCompressor* self) {

	if (self->cctx) {

		ZSTD_freeCCtx(self->cctx);

		self->cctx = NULL;

	}

	if (self->params) {

		ZSTD_freeCCtxParams(self->params);

		self->params = NULL;

	}

	Py_XDECREF(self->dict);

	PyObject_Del(self);

}

PyDoc_STRVAR(ZstdCompressor_memory_size__doc__,

"memory_size()\n"

"\n"

"Obtain the memory usage of this compressor, in bytes.\n"

);

static PyObject* ZstdCompressor_memory_size(ZstdCompressor* self) {

	if (self->cctx) {

		return PyLong_FromSize_t(ZSTD_sizeof_CCtx(self->cctx));

	}

	else {

		PyErr_SetString(ZstdError, "no compressor context found; this should never happen");

		return NULL;

	}

}

PyDoc_STRVAR(ZstdCompressor_frame_progression__doc__,

"frame_progression()\n"

"\n"

"Return information on how much work the compressor has done.\n"

"\n"

"Returns a 3-tuple of (ingested, consumed, produced).\n"

);

static PyObject* ZstdCompressor_frame_progression(ZstdCompressor* self) {

	return frame_progression(self->cctx);

}

PyDoc_STRVAR(ZstdCompressor_copy_stream__doc__,

"copy_stream(ifh, ofh[, size=0, read_size=default, write_size=default])\n"

"compress data between streams\n"

"\n"

"Data will be read from ``ifh``, compressed, and written to ``ofh``.\n"

"``ifh`` must have a ``read(size)`` method. ``ofh`` must have a ``write(data)``\n"

"method.\n"

"\n"

"An optional ``size`` argument specifies the size of the source stream.\n"

"If defined, compression parameters will be tuned based on the size.\n"

"\n"

"Optional arguments ``read_size`` and ``write_size`` define the chunk sizes\n"

"of ``read()`` and ``write()`` operations, respectively. By default, they use\n"

"the default compression stream input and output sizes, respectively.\n"

);

static PyObject* ZstdCompressor_copy_stream(ZstdCompressor* self, PyObject* args, PyObject* kwargs) {

	static char* kwlist[] = {

		"ifh",

		"ofh",

		"size",

		"read_size",

		"write_size",

		NULL

	};

	PyObject* source;

	PyObject* dest;

	unsigned long long sourceSize = ZSTD_CONTENTSIZE_UNKNOWN;

	size_t inSize = ZSTD_CStreamInSize();

	size_t outSize = ZSTD_CStreamOutSize();

	ZSTD_inBuffer input;

	ZSTD_outBuffer output;

	Py_ssize_t totalRead = 0;

	Py_ssize_t totalWrite = 0;

	char* readBuffer;

	Py_ssize_t readSize;

	PyObject* readResult = NULL;

	PyObject* res = NULL;

	size_t zresult;

	PyObject* writeResult;

	PyObject* totalReadPy;

	PyObject* totalWritePy;

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO|Kkk:copy_stream", kwlist,

		&source, &dest, &sourceSize, &inSize, &outSize)) {

		return NULL;

	}

	if (!PyObject_HasAttrString(source, "read")) {

		PyErr_SetString(PyExc_ValueError, "first argument must have a read() method");

		return NULL;

	}

	if (!PyObject_HasAttrString(dest, "write")) {

		PyErr_SetString(PyExc_ValueError, "second argument must have a write() method");

		return NULL;

	}

	if (ensure_cctx(self)) {

		return NULL;

	}

	zresult = ZSTD_CCtx_setPledgedSrcSize(self->cctx, sourceSize);

	if (ZSTD_isError(zresult)) {

		PyErr_Format(ZstdError, "error setting source size: %s",

			ZSTD_getErrorName(zresult));

		return NULL;

	}

	/* Prevent free on uninitialized memory in finally. */

	output.dst = PyMem_Malloc(outSize);

	if (!output.dst) {

		PyErr_NoMemory();

		res = NULL;

		goto finally;

	}

	output.size = outSize;

	output.pos = 0;

	input.src = NULL;

	input.size = 0;

	input.pos = 0;

	while (1) {

		/* Try to read from source stream. */

		readResult = PyObject_CallMethod(source, "read", "n", inSize);

		if (!readResult) {

			PyErr_SetString(ZstdError, "could not read() from source");

			goto finally;

		}

		PyBytes_AsStringAndSize(readResult, &readBuffer, &readSize);

		/* If no data was read, we're at EOF. */

		if (0 == readSize) {

			break;

		}

		totalRead += readSize;

		/* Send data to compressor */

		input.src = readBuffer;

		input.size = readSize;

		input.pos = 0;

		while (input.pos < input.size) {

			Py_BEGIN_ALLOW_THREADS

			zresult = ZSTD_compress_generic(self->cctx, &output, &input, ZSTD_e_continue);

			Py_END_ALLOW_THREADS

			if (ZSTD_isError(zresult)) {

				res = NULL;

				PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult));

				goto finally;

			}

			if (output.pos) {

#if PY_MAJOR_VERSION >= 3

				writeResult = PyObject_CallMethod(dest, "write", "y#",

#else

				writeResult = PyObject_CallMethod(dest, "write", "s#",

#endif

					output.dst, output.pos);

				Py_XDECREF(writeResult);

				totalWrite += output.pos;

				output.pos = 0;

			}

		}

		Py_CLEAR(readResult);

	}

	/* We've finished reading. Now flush the compressor stream. */

	assert(input.pos == input.size);

	while (1) {

		Py_BEGIN_ALLOW_THREADS

		zresult = ZSTD_compress_generic(self->cctx, &output, &input, ZSTD_e_end);

		Py_END_ALLOW_THREADS

		if (ZSTD_isError(zresult)) {

			PyErr_Format(ZstdError, "error ending compression stream: %s",

				ZSTD_getErrorName(zresult));

			res = NULL;

			goto finally;

		}

		if (output.pos) {

#if PY_MAJOR_VERSION >= 3

			writeResult = PyObject_CallMethod(dest, "write", "y#",

#else

			writeResult = PyObject_CallMethod(dest, "write", "s#",

#endif

				output.dst, output.pos);

			totalWrite += output.pos;

			Py_XDECREF(writeResult);

			output.pos = 0;

		}

		if (!zresult) {

			break;

		}

	}

	totalReadPy = PyLong_FromSsize_t(totalRead);

	totalWritePy = PyLong_FromSsize_t(totalWrite);

	res = PyTuple_Pack(2, totalReadPy, totalWritePy);

	Py_DECREF(totalReadPy);

	Py_DECREF(totalWritePy);

finally:

	if (output.dst) {

		PyMem_Free(output.dst);

	}

	Py_XDECREF(readResult);

	return res;

}

PyDoc_STRVAR(ZstdCompressor_stream_reader__doc__,

"stream_reader(source, [size=0])\n"

"\n"

"Obtain an object that behaves like an I/O stream.\n"

"\n"

"The source object can be any object with a ``read(size)`` method\n"

"or an object that conforms to the buffer protocol.\n"

);

static ZstdCompressionReader* ZstdCompressor_stream_reader(ZstdCompressor* self, PyObject* args, PyObject* kwargs) {

	static char* kwlist[] = {

		"source",

		"size",

		"read_size",

		NULL

	};

	PyObject* source;

	unsigned long long sourceSize = ZSTD_CONTENTSIZE_UNKNOWN;

	size_t readSize = ZSTD_CStreamInSize();

	ZstdCompressionReader* result = NULL;

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|Kk:stream_reader", kwlist,

		&source, &sourceSize, &readSize)) {

		return NULL;

	}

	result = (ZstdCompressionReader*)PyObject_CallObject((PyObject*)&ZstdCompressionReaderType, NULL);

	if (!result) {

		return NULL;

	}

	if (PyObject_HasAttrString(source, "read")) {

		result->reader = source;

		Py_INCREF(source);