/**

* 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(ZstdCompressionReader* self) {

	Py_XDECREF(self->compressor);

	Py_XDECREF(self->reader);

	if (self->buffer.buf) {

		PyBuffer_Release(&self->buffer);

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

	}

	PyObject_Del(self);

}

static ZstdCompressionReader* reader_enter(ZstdCompressionReader* 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(ZstdCompressionReader* 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 associated with source. */

	Py_CLEAR(self->reader);

	if (self->buffer.buf) {

		PyBuffer_Release(&self->buffer);

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

	}

    Py_CLEAR(self->compressor);

	Py_RETURN_FALSE;

}

static PyObject* reader_readable(ZstdCompressionReader* self) {

	Py_RETURN_TRUE;

}

static PyObject* reader_writable(ZstdCompressionReader* self) {

	Py_RETURN_FALSE;

}

static PyObject* reader_seekable(ZstdCompressionReader* self) {

	Py_RETURN_FALSE;

}

static PyObject* reader_readline(PyObject* self, PyObject* args) {

	set_unsupported_operation();

	return NULL;

}

static PyObject* reader_readlines(PyObject* self, PyObject* args) {

	set_unsupported_operation();

	return NULL;

}

static PyObject* reader_write(PyObject* self, PyObject* args) {

	PyErr_SetString(PyExc_OSError, "stream is not writable");

	return NULL;

}

static PyObject* reader_writelines(PyObject* self, PyObject* args) {

	PyErr_SetString(PyExc_OSError, "stream is not writable");

	return NULL;

}

static PyObject* reader_isatty(PyObject* self) {

	Py_RETURN_FALSE;

}

static PyObject* reader_flush(PyObject* self) {

	Py_RETURN_NONE;

}

static PyObject* reader_close(ZstdCompressionReader* self) {

	self->closed = 1;

	Py_RETURN_NONE;

}

static PyObject* reader_tell(ZstdCompressionReader* self) {

	/* TODO should this raise OSError since stream isn't seekable? */

	return PyLong_FromUnsignedLongLong(self->bytesCompressed);

}

int read_compressor_input(ZstdCompressionReader* 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);

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

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

		self->input.pos = 0;

	}

	return 1;

}

int compress_input(ZstdCompressionReader* self, ZSTD_outBuffer* output) {

	size_t oldPos;

	size_t zresult;

	/* If we have data left over, consume it. */

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

		oldPos = output->pos;

		Py_BEGIN_ALLOW_THREADS

		zresult = ZSTD_compressStream2(self->compressor->cctx,

		    output, &self->input, ZSTD_e_continue);

		Py_END_ALLOW_THREADS

		self->bytesCompressed += output->pos - oldPos;

		/* Input exhausted. Clear out 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 compress error: %s", ZSTD_getErrorName(zresult));

			return -1;

		}

	}

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

        return 1;

    }

    else {

        return 0;

    }

}

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

	static char* kwlist[] = {

		"size",

		NULL

	};

	Py_ssize_t size = -1;

	PyObject* result = NULL;

	char* resultBuffer;

	Py_ssize_t resultSize;

	size_t zresult;

	size_t oldPos;

	int readResult, compressResult;

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

	self->output.dst = resultBuffer;

	self->output.size = resultSize;

	self->output.pos = 0;

readinput:

    compressResult = compress_input(self, &self->output);

	if (-1 == compressResult) {

		Py_XDECREF(result);

		return NULL;

	}

	else if (0 == compressResult) {

		/* There is room in the output. We fall through to below, which will

		 * either get more input for us or will attempt to end the stream.

		 */

	}

	else if (1 == compressResult) {

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

		return result;

	}

	else {

		assert(0);

	}

	readResult = read_compressor_input(self);

	if (-1 == readResult) {

		return NULL;

	}

	else if (0 == readResult) { }

	else if (1 == readResult) { }

	else {

		assert(0);

	}

	if (self->input.size) {

		goto readinput;

	}

	/* Else EOF */

	oldPos = self->output.pos;

	zresult = ZSTD_compressStream2(self->compressor->cctx, &self->output,

		&self->input, ZSTD_e_end);

	self->bytesCompressed += self->output.pos - oldPos;

	if (ZSTD_isError(zresult)) {

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

			ZSTD_getErrorName(zresult));

		Py_XDECREF(result);

		return NULL;

	}

	assert(self->output.pos);

	if (0 == zresult) {

		self->finishedOutput = 1;

	}

	if (safe_pybytes_resize(&result, self->output.pos)) {

		Py_XDECREF(result);

		return NULL;

	}

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

	return result;

}

static PyObject* reader_read1(ZstdCompressionReader* 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 compressResult;

	size_t oldPos;

	size_t zresult;

	if (self->closed) {

		PyErr_SetString(PyExc_ValueError, "stream is closed");

		return NULL;

	}

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|n:read1", 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_CStreamOutSize();

	}

	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 compression because

	   not every input will generate output. We /could/ flush the compressor,

	   but this may not be desirable. We allow multiple read() from the

	   underlying stream. But unlike read(), we return as soon as output data

	   is available.

	*/

	compressResult = compress_input(self, &output);

	if (-1 == compressResult) {

		Py_XDECREF(result);

		return NULL;

	}

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

	else {

		assert(0);

	}

	if (output.pos) {

		goto finally;

	}

	while (!self->finishedInput) {

		int readResult = read_compressor_input(self);