/**

* 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;

/**

 * Ensure the ZSTD_DCtx on a decompressor is initiated and ready for a new operation.

 */

int ensure_dctx(ZstdDecompressor* decompressor, int loadDict) {

	size_t zresult;

	ZSTD_DCtx_reset(decompressor->dctx, ZSTD_reset_session_only);

	if (decompressor->maxWindowSize) {

		zresult = ZSTD_DCtx_setMaxWindowSize(decompressor->dctx, decompressor->maxWindowSize);

		if (ZSTD_isError(zresult)) {

			PyErr_Format(ZstdError, "unable to set max window size: %s",

				ZSTD_getErrorName(zresult));

			return 1;

		}

	}

	zresult = ZSTD_DCtx_setFormat(decompressor->dctx, decompressor->format);

	if (ZSTD_isError(zresult)) {

		PyErr_Format(ZstdError, "unable to set decoding format: %s",

			ZSTD_getErrorName(zresult));

		return 1;

	}

	if (loadDict && decompressor->dict) {

		if (ensure_ddict(decompressor->dict)) {

			return 1;

		}

		zresult = ZSTD_DCtx_refDDict(decompressor->dctx, decompressor->dict->ddict);

		if (ZSTD_isError(zresult)) {

			PyErr_Format(ZstdError, "unable to reference prepared dictionary: %s",

				ZSTD_getErrorName(zresult));

			return 1;

		}

	}

	return 0;

}

PyDoc_STRVAR(Decompressor__doc__,

"ZstdDecompressor(dict_data=None)\n"

"\n"

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

"\n"

"An instance can perform multiple decompression operations."

);

static int Decompressor_init(ZstdDecompressor* self, PyObject* args, PyObject* kwargs) {

	static char* kwlist[] = {

		"dict_data",

		"max_window_size",

		"format",

		NULL

	};

	ZstdCompressionDict* dict = NULL;

	ZSTD_format_e format = ZSTD_f_zstd1;

	self->dctx = NULL;

	self->dict = NULL;

		&ZstdCompressionDictType, &dict, &maxWindowSize, &format)) {

		return -1;

	}

	self->dctx = ZSTD_createDCtx();

	if (!self->dctx) {

		PyErr_NoMemory();

		goto except;

	}

	self->maxWindowSize = maxWindowSize;

	self->format = format;

	if (dict) {

		self->dict = dict;

		Py_INCREF(dict);

	}

	if (ensure_dctx(self, 1)) {

		goto except;

	}

	return 0;

except:

	Py_CLEAR(self->dict);

	if (self->dctx) {

		ZSTD_freeDCtx(self->dctx);

		self->dctx = NULL;

	}

	return -1;

}

static void Decompressor_dealloc(ZstdDecompressor* self) {

	Py_CLEAR(self->dict);

	if (self->dctx) {

		ZSTD_freeDCtx(self->dctx);

		self->dctx = NULL;

	}

	PyObject_Del(self);

}

PyDoc_STRVAR(Decompressor_memory_size__doc__,

"memory_size() -- Size of decompression context, in bytes\n"

);

static PyObject* Decompressor_memory_size(ZstdDecompressor* self) {

	if (self->dctx) {

		return PyLong_FromSize_t(ZSTD_sizeof_DCtx(self->dctx));

	}

	else {

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

		return NULL;

	}

}

PyDoc_STRVAR(Decompressor_copy_stream__doc__,

	"copy_stream(ifh, ofh[, read_size=default, write_size=default]) -- decompress data between streams\n"

	"\n"

	"Compressed data will be read from ``ifh``, decompressed, and written to\n"

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

	"``write(data)`` method.\n"

	"\n"

	"The optional ``read_size`` and ``write_size`` arguments control the chunk\n"

	"size of data that is ``read()`` and ``write()`` between streams. They default\n"

	"to the default input and output sizes of zstd decompressor streams.\n"

);

static PyObject* Decompressor_copy_stream(ZstdDecompressor* self, PyObject* args, PyObject* kwargs) {

	static char* kwlist[] = {

		"ifh",

		"ofh",

		"read_size",

		"write_size",

		NULL

	};

	PyObject* source;

	PyObject* dest;

	size_t inSize = ZSTD_DStreamInSize();

	size_t outSize = ZSTD_DStreamOutSize();

	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 = 0;

	PyObject* writeResult;

	PyObject* totalReadPy;

	PyObject* totalWritePy;

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

		&source, &dest, &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;

	}

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

	output.dst = NULL;

	if (ensure_dctx(self, 1)) {

		res = NULL;

		goto finally;

	}

	output.dst = PyMem_Malloc(outSize);

	if (!output.dst) {

		PyErr_NoMemory();

		res = NULL;

		goto finally;

	}

	output.size = outSize;

	output.pos = 0;

	/* Read source stream until EOF */

	while (1) {

		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 decompressor */

		input.src = readBuffer;

		input.size = readSize;

		input.pos = 0;

		while (input.pos < input.size) {

			Py_BEGIN_ALLOW_THREADS

			zresult = ZSTD_decompressStream(self->dctx, &output, &input);

			Py_END_ALLOW_THREADS

			if (ZSTD_isError(zresult)) {

				PyErr_Format(ZstdError, "zstd decompressor error: %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);

				Py_XDECREF(writeResult);

				totalWrite += output.pos;

				output.pos = 0;

			}

		}

		Py_CLEAR(readResult);

	}

	/* Source stream is exhausted. Finish up. */

	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(Decompressor_decompress__doc__,

"decompress(data[, max_output_size=None]) -- Decompress data in its entirety\n"

"\n"

"This method will decompress the entirety of the argument and return the\n"

"result.\n"

"\n"

"The input bytes are expected to contain a full Zstandard frame (something\n"

"compressed with ``ZstdCompressor.compress()`` or similar). If the input does\n"

"not contain a full frame, an exception will be raised.\n"

"\n"

"If the frame header of the compressed data does not contain the content size\n"

"``max_output_size`` must be specified or ``ZstdError`` will be raised. An\n"

"allocation of size ``max_output_size`` will be performed and an attempt will\n"

"be made to perform decompression into that buffer. If the buffer is too\n"

"small or cannot be allocated, ``ZstdError`` will be raised. The buffer will\n"

"be resized if it is too large.\n"

"\n"

"Uncompressed data could be much larger than compressed data. As a result,\n"

"calling this function could result in a very large memory allocation being\n"

"performed to hold the uncompressed data. Therefore it is **highly**\n"

"recommended to use a streaming decompression method instead of this one.\n"

);

PyObject* Decompressor_decompress(ZstdDecompressor* self, PyObject* args, PyObject* kwargs) {

	static char* kwlist[] = {

		"data",

		"max_output_size",

		NULL

	};

	Py_buffer source;

	Py_ssize_t maxOutputSize = 0;

	unsigned long long decompressedSize;

	size_t destCapacity;

	PyObject* result = NULL;

	size_t zresult;

	ZSTD_outBuffer outBuffer;

	ZSTD_inBuffer inBuffer;

#if PY_MAJOR_VERSION >= 3

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*|n:decompress",

#else

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*|n:decompress",

#endif

		kwlist, &source, &maxOutputSize)) {

		return NULL;

	}

	if (!PyBuffer_IsContiguous(&source, 'C') || source.ndim > 1) {

		PyErr_SetString(PyExc_ValueError,

			"data buffer should be contiguous and have at most one dimension");

		goto finally;

	}

	if (ensure_dctx(self, 1)) {

		goto finally;

	}

	decompressedSize = ZSTD_getFrameContentSize(source.buf, source.len);

	if (ZSTD_CONTENTSIZE_ERROR == decompressedSize) {

		PyErr_SetString(ZstdError, "error determining content size from frame header");

		goto finally;

	}

	/* Special case of empty frame. */

	else if (0 == decompressedSize) {

		result = PyBytes_FromStringAndSize("", 0);

		goto finally;

	}

	/* Missing content size in frame header. */

	if (ZSTD_CONTENTSIZE_UNKNOWN == decompressedSize) {

		if (0 == maxOutputSize) {

			PyErr_SetString(ZstdError, "could not determine content size in frame header");

			goto finally;

		}

		result = PyBytes_FromStringAndSize(NULL, maxOutputSize);

		destCapacity = maxOutputSize;

		decompressedSize = 0;

	}

	/* Size is recorded in frame header. */

	else {

		assert(SIZE_MAX >= PY_SSIZE_T_MAX);

		if (decompressedSize > PY_SSIZE_T_MAX) {

			PyErr_SetString(ZstdError, "frame is too large to decompress on this platform");

			goto finally;

		}

		result = PyBytes_FromStringAndSize(NULL, (Py_ssize_t)decompressedSize);

		destCapacity = (size_t)decompressedSize;

	}

	if (!result) {

		goto finally;

	}

	outBuffer.dst = PyBytes_AsString(result);

	outBuffer.size = destCapacity;

	outBuffer.pos = 0;

	inBuffer.src = source.buf;

	inBuffer.size = source.len;

	inBuffer.pos = 0;

	Py_BEGIN_ALLOW_THREADS

	zresult = ZSTD_decompressStream(self->dctx, &outBuffer, &inBuffer);

	Py_END_ALLOW_THREADS

	if (ZSTD_isError(zresult)) {

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

		Py_CLEAR(result);

		goto finally;

	}

	else if (zresult) {

		PyErr_Format(ZstdError, "decompression error: did not decompress full frame");

		Py_CLEAR(result);

		goto finally;

	}

	else if (decompressedSize && outBuffer.pos != decompressedSize) {

		PyErr_Format(ZstdError, "decompression error: decompressed %zu bytes; expected %llu",

			zresult, decompressedSize);

		Py_CLEAR(result);

		goto finally;

	}

	else if (outBuffer.pos < destCapacity) {

		if (safe_pybytes_resize(&result, outBuffer.pos)) {

			Py_CLEAR(result);

			goto finally;

		}

	}

finally:

	PyBuffer_Release(&source);

	return result;