/**

* Copyright (c) 2017-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;

static void set_unsupported_operation(void) {

	PyObject* iomod;

	PyObject* exc;

	iomod = PyImport_ImportModule("io");

	if (NULL == iomod) {

		return;

	}

	exc = PyObject_GetAttrString(iomod, "UnsupportedOperation");

	if (NULL == exc) {

		Py_DECREF(iomod);

		return;

	}

	PyErr_SetNone(exc);

	Py_DECREF(exc);

	Py_DECREF(iomod);

}

static void reader_dealloc(ZstdDecompressionReader* self) {

	Py_XDECREF(self->decompressor);

	Py_XDECREF(self->reader);

	if (self->buffer.buf) {

		PyBuffer_Release(&self->buffer);

	}

	PyObject_Del(self);

}

static ZstdDecompressionReader* reader_enter(ZstdDecompressionReader* self) {

	if (self->entered) {

		PyErr_SetString(PyExc_ValueError, "cannot __enter__ multiple times");

		return NULL;

	}

	self->entered = 1;

	Py_INCREF(self);

	return self;

}

static PyObject* reader_exit(ZstdDecompressionReader* self, PyObject* args) {

	PyObject* exc_type;

	PyObject* exc_value;

	PyObject* exc_tb;

	if (!PyArg_ParseTuple(args, "OOO:__exit__", &exc_type, &exc_value, &exc_tb)) {

		return NULL;

	}

	self->entered = 0;

	self->closed = 1;

	/* Release resources. */

	Py_CLEAR(self->reader);

	if (self->buffer.buf) {

		PyBuffer_Release(&self->buffer);

		memset(&self->buffer, 0, sizeof(self->buffer));

	}

	Py_CLEAR(self->decompressor);

	Py_RETURN_FALSE;

}

static PyObject* reader_readable(PyObject* self) {

	Py_RETURN_TRUE;

}

static PyObject* reader_writable(PyObject* self) {

	Py_RETURN_FALSE;

}

static PyObject* reader_seekable(PyObject* self) {

	Py_RETURN_TRUE;

}

static PyObject* reader_close(ZstdDecompressionReader* self) {

	self->closed = 1;

	Py_RETURN_NONE;

}

static PyObject* reader_flush(PyObject* self) {

	Py_RETURN_NONE;

}

static PyObject* reader_isatty(PyObject* self) {

	Py_RETURN_FALSE;

}

/**

 * Read available input.

 *

 * Returns 0 if no data was added to input.

 * Returns 1 if new input data is available.

 * Returns -1 on error and sets a Python exception as a side-effect.

 */

int read_decompressor_input(ZstdDecompressionReader* self) {

	if (self->finishedInput) {

		return 0;

	}

	if (self->input.pos != self->input.size) {

		return 0;

	}

	if (self->reader) {

        Py_buffer buffer;

        assert(self->readResult == NULL);

        self->readResult = PyObject_CallMethod(self->reader, "read",

            "k", self->readSize);

        if (NULL == self->readResult) {

            return -1;

        }

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

        if (0 != PyObject_GetBuffer(self->readResult, &buffer, PyBUF_CONTIG_RO)) {

            return -1;

        }

        /* EOF */

        if (0 == buffer.len) {

            self->finishedInput = 1;

            Py_CLEAR(self->readResult);

        }

        else {

            self->input.src = buffer.buf;

            self->input.size = buffer.len;

            self->input.pos = 0;

        }

        PyBuffer_Release(&buffer);

	}

	else {

		assert(self->buffer.buf);

        /*

         * We should only get here once since expectation is we always

         * exhaust input buffer before reading again.

         */

        assert(self->input.src == NULL);

		self->input.src = self->buffer.buf;

        self->input.size = self->buffer.len;

        self->input.pos = 0;

	}

	return 1;

}

/**

 * Decompresses available input into an output buffer.

 *

 * Returns 0 if we need more input.

 * Returns 1 if output buffer should be emitted.

 * Returns -1 on error and sets a Python exception.

 */

int decompress_input(ZstdDecompressionReader* self, ZSTD_outBuffer* output) {

	size_t zresult;

	if (self->input.pos >= self->input.size) {

		return 0;

	}

	Py_BEGIN_ALLOW_THREADS

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

	Py_END_ALLOW_THREADS

	/* Input exhausted. Clear our state tracking. */

	if (self->input.pos == self->input.size) {

		memset(&self->input, 0, sizeof(self->input));

		Py_CLEAR(self->readResult);

		if (self->buffer.buf) {

			self->finishedInput = 1;

		}

	}

	if (ZSTD_isError(zresult)) {

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

		return -1;

	}

	/* We fulfilled the full read request. Signal to emit. */

	if (output->pos && output->pos == output->size) {

		return 1;

	}

	/* We're at the end of a frame and we aren't allowed to return data

	   spanning frames. */

	else if (output->pos && zresult == 0 && !self->readAcrossFrames) {

		return 1;

	}

	/* There is more room in the output. Signal to collect more data. */

	return 0;

}

static PyObject* reader_read(ZstdDecompressionReader* self, PyObject* args, PyObject* kwargs) {

	static char* kwlist[] = {

		"size",

		NULL

	};

	Py_ssize_t size = -1;

	PyObject* result = NULL;

	char* resultBuffer;

	Py_ssize_t resultSize;

	ZSTD_outBuffer output;

	int decompressResult, readResult;

	if (self->closed) {

		PyErr_SetString(PyExc_ValueError, "stream is closed");

		return NULL;

	}

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|n", kwlist, &size)) {

		return NULL;

	}

	if (size < -1) {

		PyErr_SetString(PyExc_ValueError, "cannot read negative amounts less than -1");

		return NULL;

	}

	if (size == -1) {

		return PyObject_CallMethod((PyObject*)self, "readall", NULL);

	}

	if (self->finishedOutput || size == 0) {

		return PyBytes_FromStringAndSize("", 0);

	}

	result = PyBytes_FromStringAndSize(NULL, size);

	if (NULL == result) {

		return NULL;

	}

	PyBytes_AsStringAndSize(result, &resultBuffer, &resultSize);

	output.dst = resultBuffer;

	output.size = resultSize;

	output.pos = 0;

readinput:

	decompressResult = decompress_input(self, &output);

	if (-1 == decompressResult) {

		Py_XDECREF(result);

		return NULL;

	}

	else if (0 == decompressResult) { }

	else if (1 == decompressResult) {

		self->bytesDecompressed += output.pos;

		if (output.pos != output.size) {

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

				Py_XDECREF(result);

				return NULL;

			}

		}

		return result;

	}

	else {

		assert(0);

	}

	readResult = read_decompressor_input(self);

	if (-1 == readResult) {

		Py_XDECREF(result);

		return NULL;

	}

	else if (0 == readResult) {}

	else if (1 == readResult) {}

	else {

		assert(0);

	}

	if (self->input.size) {

		goto readinput;

	}

	/* EOF */

	self->bytesDecompressed += output.pos;

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

		Py_XDECREF(result);

		return NULL;

	}

	return result;

}

static PyObject* reader_read1(ZstdDecompressionReader* self, PyObject* args, PyObject* kwargs) {

	static char* kwlist[] = {

		"size",

		NULL

	};

	Py_ssize_t size = -1;

	PyObject* result = NULL;

	char* resultBuffer;

	Py_ssize_t resultSize;

	ZSTD_outBuffer output;

	if (self->closed) {

		PyErr_SetString(PyExc_ValueError, "stream is closed");

		return NULL;

	}

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|n", kwlist, &size)) {

		return NULL;

	}

	if (size < -1) {

		PyErr_SetString(PyExc_ValueError, "cannot read negative amounts less than -1");

		return NULL;

	}

	if (self->finishedOutput || size == 0) {

		return PyBytes_FromStringAndSize("", 0);

	}

	if (size == -1) {

		size = ZSTD_DStreamOutSize();

	}

	result = PyBytes_FromStringAndSize(NULL, size);

	if (NULL == result) {

		return NULL;

	}

	PyBytes_AsStringAndSize(result, &resultBuffer, &resultSize);

	output.dst = resultBuffer;

	output.size = resultSize;

	output.pos = 0;

	/* read1() is supposed to use at most 1 read() from the underlying stream.

	 * However, we can't satisfy this requirement with decompression due to the

	 * nature of how decompression works. Our strategy is to read + decompress

	 * until we get any output, at which point we return. This satisfies the

	 * intent of the read1() API to limit read operations.

	 */

	while (!self->finishedInput) {

		int readResult, decompressResult;

		readResult = read_decompressor_input(self);

		if (-1 == readResult) {

			Py_XDECREF(result);

			return NULL;

		}

		else if (0 == readResult || 1 == readResult) { }

		else {

			assert(0);

		}

		decompressResult = decompress_input(self, &output);

		if (-1 == decompressResult) {

			Py_XDECREF(result);

			return NULL;

		}

		else if (0 == decompressResult || 1 == decompressResult) { }

		else {

			assert(0);

		}

		if (output.pos) {

		    break;

		}

	}

	self->bytesDecompressed += output.pos;

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

		Py_XDECREF(result);

		return NULL;

	}

	return result;

}

static PyObject* reader_readinto(ZstdDecompressionReader* self, PyObject* args) {

	Py_buffer dest;

	ZSTD_outBuffer output;

	int decompressResult, readResult;

	PyObject* result = NULL;

	if (self->closed) {

		PyErr_SetString(PyExc_ValueError, "stream is closed");

		return NULL;

	}