/**

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

extern PyObject* ZstdError;

ZSTD_DStream* DStream_from_ZstdDecompressor(ZstdDecompressor* decompressor) {

	ZSTD_DStream* dstream;

	void* dictData = NULL;

	size_t dictSize = 0;

	size_t zresult;

	dstream = ZSTD_createDStream();

	if (!dstream) {

		PyErr_SetString(ZstdError, "could not create DStream");

		return NULL;

	}

	if (decompressor->dict) {

		dictData = decompressor->dict->dictData;

		dictSize = decompressor->dict->dictSize;

	}

	if (dictData) {

		zresult = ZSTD_initDStream_usingDict(dstream, dictData, dictSize);

	}

	else {

		zresult = ZSTD_initDStream(dstream);

	}

	if (ZSTD_isError(zresult)) {

		PyErr_Format(ZstdError, "could not initialize DStream: %s",

			ZSTD_getErrorName(zresult));

		return NULL;

	}

	return dstream;

}

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",

		NULL

	};

	ZstdCompressionDict* dict = NULL;

	self->refdctx = NULL;

	self->dict = NULL;

	self->ddict = NULL;

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O!", kwlist,

		&ZstdCompressionDictType, &dict)) {

		return -1;

	}

	/* Instead of creating a ZSTD_DCtx for every decompression operation,

	   we create an instance at object creation time and recycle it via

	   ZSTD_copyDCTx() on each use. This means each use is a malloc+memcpy

	   instead of a malloc+init. */

	/* TODO lazily initialize the reference ZSTD_DCtx on first use since

	   not instances of ZstdDecompressor will use a ZSTD_DCtx. */

	self->refdctx = ZSTD_createDCtx();

	if (!self->refdctx) {

		PyErr_NoMemory();

		goto except;

	}

	if (dict) {

		self->dict = dict;

		Py_INCREF(dict);

	}

	return 0;

except:

	if (self->refdctx) {

		ZSTD_freeDCtx(self->refdctx);

		self->refdctx = NULL;

	}

	return -1;

}

static void Decompressor_dealloc(ZstdDecompressor* self) {

	if (self->refdctx) {

		ZSTD_freeDCtx(self->refdctx);

	}

	Py_XDECREF(self->dict);

	if (self->ddict) {

		ZSTD_freeDDict(self->ddict);

		self->ddict = NULL;

	}

	PyObject_Del(self);

}

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

	ZSTD_inBuffer input;

	ZSTD_outBuffer output;

	Py_ssize_t totalRead = 0;

	Py_ssize_t totalWrite = 0;

	char* readBuffer;

	Py_ssize_t readSize;

	PyObject* readResult;

	PyObject* res = NULL;

	size_t zresult = 0;

	PyObject* writeResult;

	PyObject* totalReadPy;

	PyObject* totalWritePy;

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO|kk", 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;

	}

	dstream = DStream_from_ZstdDecompressor(self);

	if (!dstream) {

		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(dstream, &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;

			}

		}

	}

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

	ZSTD_freeDStream(dstream);

	dstream = NULL;

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

	}

	if (dstream) {

		ZSTD_freeDStream(dstream);

	}

	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

	};

	const char* source;

	Py_ssize_t sourceSize;

	Py_ssize_t maxOutputSize = 0;

	unsigned long long decompressedSize;

	size_t destCapacity;

	PyObject* result = NULL;

	ZSTD_DCtx* dctx = NULL;

	void* dictData = NULL;

	size_t dictSize = 0;

	size_t zresult;

#if PY_MAJOR_VERSION >= 3

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y#|n", kwlist,

#else

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s#|n", kwlist,

#endif

		&source, &sourceSize, &maxOutputSize)) {

		return NULL;

	}

	dctx = PyMem_Malloc(ZSTD_sizeof_DCtx(self->refdctx));

	if (!dctx) {

		PyErr_NoMemory();

		return NULL;

	}

	ZSTD_copyDCtx(dctx, self->refdctx);

	if (self->dict) {

		dictData = self->dict->dictData;

		dictSize = self->dict->dictSize;

	}

	if (dictData && !self->ddict) {

		Py_BEGIN_ALLOW_THREADS

		self->ddict = ZSTD_createDDict(dictData, dictSize);

		Py_END_ALLOW_THREADS

		if (!self->ddict) {

			PyErr_SetString(ZstdError, "could not create decompression dict");

			goto except;

		}

	}

	decompressedSize = ZSTD_getDecompressedSize(source, sourceSize);

	/* 0 returned if content size not in the zstd frame header */

	if (0 == decompressedSize) {

		if (0 == maxOutputSize) {

			PyErr_SetString(ZstdError, "input data invalid or missing content size "

				"in frame header");

			goto except;

		}

		else {

			result = PyBytes_FromStringAndSize(NULL, maxOutputSize);

			destCapacity = maxOutputSize;

		}

	}

	else {

		result = PyBytes_FromStringAndSize(NULL, decompressedSize);

		destCapacity = decompressedSize;

	}

	if (!result) {

		goto except;

	}

	Py_BEGIN_ALLOW_THREADS

	if (self->ddict) {

		zresult = ZSTD_decompress_usingDDict(dctx, PyBytes_AsString(result), destCapacity,

			source, sourceSize, self->ddict);

	}

	else {

		zresult = ZSTD_decompressDCtx(dctx, PyBytes_AsString(result), destCapacity, source, sourceSize);

	}

	Py_END_ALLOW_THREADS

	if (ZSTD_isError(zresult)) {

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

		goto except;

	}

	else if (decompressedSize && zresult != decompressedSize) {

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

			zresult, decompressedSize);

		goto except;

	}

	else if (zresult < destCapacity) {

		if (_PyBytes_Resize(&result, zresult)) {

			goto except;

		}

	}

	goto finally;

except:

	Py_DecRef(result);

	result = NULL;

finally:

	if (dctx) {

		PyMem_FREE(dctx);

	}

	return result;

}

PyDoc_STRVAR(Decompressor_decompressobj__doc__,

"decompressobj()\n"

"\n"

"Incrementally feed data into a decompressor.\n"

"\n"

"The returned object exposes a ``decompress(data)`` method. This makes it\n"

"compatible with ``zlib.decompressobj`` and ``bz2.BZ2Decompressor`` so that\n"

"callers can swap in the zstd decompressor while using the same API.\n"

);

static ZstdDecompressionObj* Decompressor_decompressobj(ZstdDecompressor* self) {

	ZstdDecompressionObj* result = PyObject_New(ZstdDecompressionObj, &ZstdDecompressionObjType);

	if (!result) {

		return NULL;

	}

	result->dstream = DStream_from_ZstdDecompressor(self);

	if (!result->dstream) {

		Py_DecRef((PyObject*)result);

		return NULL;

	}

	result->decompressor = self;

	Py_INCREF(result->decompressor);

	result->finished = 0;

	return result;

}

PyDoc_STRVAR(Decompressor_read_from__doc__,

"read_from(reader[, read_size=default, write_size=default, skip_bytes=0])\n"

"Read compressed data and return an iterator\n"

"\n"

"Returns an iterator of decompressed data chunks produced from reading from\n"

"the ``reader``.\n"

"\n"

"Compressed data will be obtained from ``reader`` by calling the\n"

"``read(size)`` method of it. The source data will be streamed into a\n"

"decompressor. As decompressed data is available, it will be exposed to the\n"

"returned iterator.\n"

"\n"

"Data is ``read()`` in chunks of size ``read_size`` and exposed to the\n"

"iterator in chunks of size ``write_size``. The default values are the input\n"

"and output sizes for a zstd streaming decompressor.\n"

"\n"

"There is also support for skipping the first ``skip_bytes`` of data from\n"

"the source.\n"

);

static ZstdDecompressorIterator* Decompressor_read_from(ZstdDecompressor* self, PyObject* args, PyObject* kwargs) {

	static char* kwlist[] = {

		"reader",

		"read_size",

		"write_size",

		"skip_bytes",

		NULL

	};

	PyObject* reader;

	size_t inSize = ZSTD_DStreamInSize();

	size_t outSize = ZSTD_DStreamOutSize();

	ZstdDecompressorIterator* result;

	size_t skipBytes = 0;

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|kkk", kwlist, &reader,

		&inSize, &outSize, &skipBytes)) {

		return NULL;

	}

	if (skipBytes >= inSize) {

		PyErr_SetString(PyExc_ValueError,

			"skip_bytes must be smaller than read_size");

		return NULL;

	}

	result = PyObject_New(ZstdDecompressorIterator, &ZstdDecompressorIteratorType);

	if (!result) {

		return NULL;

	}

	result->decompressor = NULL;

	result->reader = NULL;

	result->buffer = NULL;

	result->dstream = NULL;

	result->input.src = NULL;

	result->output.dst = NULL;

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

		result->reader = reader;

		Py_INCREF(result->reader);

	}

	else if (1 == PyObject_CheckBuffer(reader)) {

		/* Object claims it is a buffer. Try to get a handle to it. */

		result->buffer = PyMem_Malloc(sizeof(Py_buffer));

		if (!result->buffer) {

			goto except;

		}

		memset(result->buffer, 0, sizeof(Py_buffer));

		if (0 != PyObject_GetBuffer(reader, result->buffer, PyBUF_CONTIG_RO)) {

			goto except;

		}

		result->bufferOffset = 0;

	}

	else {

		PyErr_SetString(PyExc_ValueError,

			"must pass an object with a read() method or conforms to buffer protocol");

		goto except;

	}

	result->decompressor = self;

	Py_INCREF(result->decompressor);

	result->inSize = inSize;

	result->outSize = outSize;

	result->skipBytes = skipBytes;

	result->dstream = DStream_from_ZstdDecompressor(self);

	if (!result->dstream) {

		goto except;

	}

	result->input.src = PyMem_Malloc(inSize);

	if (!result->input.src) {

		PyErr_NoMemory();

		goto except;

	}

	result->input.size = 0;

	result->input.pos = 0;

	result->output.dst = NULL;

	result->output.size = 0;

	result->output.pos = 0;

	result->readCount = 0;

	result->finishedInput = 0;