/**

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

PyDoc_STRVAR(ZstdDecompressionWriter__doc,

"""A context manager used for writing decompressed output.\n"

);

static void ZstdDecompressionWriter_dealloc(ZstdDecompressionWriter* self) {

	Py_XDECREF(self->decompressor);

	Py_XDECREF(self->writer);

	PyObject_Del(self);

}

static PyObject* ZstdDecompressionWriter_enter(ZstdDecompressionWriter* self) {

	if (self->closed) {

		PyErr_SetString(PyExc_ValueError, "stream is closed");

		return NULL;

	}

	if (self->entered) {

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

		return NULL;

	}

	self->entered = 1;

	Py_INCREF(self);

	return (PyObject*)self;

}

static PyObject* ZstdDecompressionWriter_exit(ZstdDecompressionWriter* self, PyObject* args) {

	self->entered = 0;

	if (NULL == PyObject_CallMethod((PyObject*)self, "close", NULL)) {

		return NULL;

	}

	Py_RETURN_FALSE;

}

static PyObject* ZstdDecompressionWriter_memory_size(ZstdDecompressionWriter* self) {

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

}

static PyObject* ZstdDecompressionWriter_write(ZstdDecompressionWriter* self, PyObject* args, PyObject* kwargs) {

	static char* kwlist[] = {

		"data",

		NULL

	};

	PyObject* result = NULL;

	Py_buffer source;

	size_t zresult = 0;

	ZSTD_inBuffer input;

	ZSTD_outBuffer output;

	PyObject* res;

	Py_ssize_t totalWrite = 0;

#if PY_MAJOR_VERSION >= 3

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*:write",

#else

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*:write",

#endif

		kwlist, &source)) {

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

		PyErr_SetString(PyExc_ValueError, "stream is closed");

		return NULL;

	}

	output.dst = PyMem_Malloc(self->outSize);

	if (!output.dst) {

		PyErr_NoMemory();

		goto finally;

	}

	output.size = self->outSize;

	output.pos = 0;

	input.src = source.buf;

	input.size = source.len;

	input.pos = 0;

	while (input.pos < (size_t)source.len) {

		Py_BEGIN_ALLOW_THREADS

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

		Py_END_ALLOW_THREADS

		if (ZSTD_isError(zresult)) {

			PyMem_Free(output.dst);

			PyErr_Format(ZstdError, "zstd decompress error: %s",

				ZSTD_getErrorName(zresult));

			goto finally;

		}

		if (output.pos) {

#if PY_MAJOR_VERSION >= 3

			res = PyObject_CallMethod(self->writer, "write", "y#",

#else

			res = PyObject_CallMethod(self->writer, "write", "s#",

#endif

				output.dst, output.pos);

			Py_XDECREF(res);

			totalWrite += output.pos;

			output.pos = 0;

		}

	}

	PyMem_Free(output.dst);

	if (self->writeReturnRead) {

		result = PyLong_FromSize_t(input.pos);

	}

	else {

		result = PyLong_FromSsize_t(totalWrite);

	}

finally:

	PyBuffer_Release(&source);

	return result;

}

static PyObject* ZstdDecompressionWriter_close(ZstdDecompressionWriter* self) {

	PyObject* result;

	if (self->closed) {

		Py_RETURN_NONE;

	}

	result = PyObject_CallMethod((PyObject*)self, "flush", NULL);

	self->closed = 1;

	if (NULL == result) {

		return NULL;

	}

	/* Call close on underlying stream as well. */

	if (PyObject_HasAttrString(self->writer, "close")) {

		return PyObject_CallMethod(self->writer, "close", NULL);

	}

	Py_RETURN_NONE;

}

static PyObject* ZstdDecompressionWriter_fileno(ZstdDecompressionWriter* self) {

	if (PyObject_HasAttrString(self->writer, "fileno")) {

		return PyObject_CallMethod(self->writer, "fileno", NULL);

	}

	else {

		PyErr_SetString(PyExc_OSError, "fileno not available on underlying writer");

		return NULL;

	}

}

static PyObject* ZstdDecompressionWriter_flush(ZstdDecompressionWriter* self) {

	if (self->closed) {

		PyErr_SetString(PyExc_ValueError, "stream is closed");

		return NULL;

	}

	if (PyObject_HasAttrString(self->writer, "flush")) {

		return PyObject_CallMethod(self->writer, "flush", NULL);

	}

	else {

		Py_RETURN_NONE;

	}

}

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

	Py_RETURN_FALSE;

}

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

	Py_RETURN_TRUE;

}

static PyObject* ZstdDecompressionWriter_unsupported(PyObject* self, PyObject* args, PyObject* kwargs) {

	PyObject* iomod;

	PyObject* exc;

	iomod = PyImport_ImportModule("io");

	if (NULL == iomod) {

		return NULL;

	}

	exc = PyObject_GetAttrString(iomod, "UnsupportedOperation");

	if (NULL == exc) {

		Py_DECREF(iomod);

		return NULL;

	}

	PyErr_SetNone(exc);

	Py_DECREF(exc);

	Py_DECREF(iomod);

	return NULL;

}

static PyMethodDef ZstdDecompressionWriter_methods[] = {

	{ "__enter__", (PyCFunction)ZstdDecompressionWriter_enter, METH_NOARGS,

	PyDoc_STR("Enter a decompression context.") },

	{ "__exit__", (PyCFunction)ZstdDecompressionWriter_exit, METH_VARARGS,

	PyDoc_STR("Exit a decompression context.") },

	{ "memory_size", (PyCFunction)ZstdDecompressionWriter_memory_size, METH_NOARGS,

	PyDoc_STR("Obtain the memory size in bytes of the underlying decompressor.") },

	{ "close", (PyCFunction)ZstdDecompressionWriter_close, METH_NOARGS, NULL },

	{ "fileno", (PyCFunction)ZstdDecompressionWriter_fileno, METH_NOARGS, NULL },

	{ "flush", (PyCFunction)ZstdDecompressionWriter_flush, METH_NOARGS, NULL },

	{ "isatty", ZstdDecompressionWriter_false, METH_NOARGS, NULL },

	{ "readable", ZstdDecompressionWriter_false, METH_NOARGS, NULL },

	{ "readline", (PyCFunction)ZstdDecompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },

	{ "readlines", (PyCFunction)ZstdDecompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },

	{ "seek", (PyCFunction)ZstdDecompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },

	{ "seekable", ZstdDecompressionWriter_false, METH_NOARGS, NULL },

	{ "tell", (PyCFunction)ZstdDecompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },

	{ "truncate", (PyCFunction)ZstdDecompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },

	{ "writable", ZstdDecompressionWriter_true, METH_NOARGS, NULL },

	{ "writelines" , (PyCFunction)ZstdDecompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },

	{ "read", (PyCFunction)ZstdDecompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },

	{ "readall", (PyCFunction)ZstdDecompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },

	{ "readinto", (PyCFunction)ZstdDecompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },

	{ "write", (PyCFunction)ZstdDecompressionWriter_write, METH_VARARGS | METH_KEYWORDS,

	PyDoc_STR("Compress data") },

	{ NULL, NULL }

};

static PyMemberDef ZstdDecompressionWriter_members[] = {

	{ "closed", T_BOOL, offsetof(ZstdDecompressionWriter, closed), READONLY, NULL },

	{ NULL }

};

PyTypeObject ZstdDecompressionWriterType = {

	PyVarObject_HEAD_INIT(NULL, 0)

	"zstd.ZstdDecompressionWriter", /* tp_name */

	sizeof(ZstdDecompressionWriter),/* tp_basicsize */

	0,                              /* tp_itemsize */

	(destructor)ZstdDecompressionWriter_dealloc, /* tp_dealloc */

	0,                              /* tp_print */

	0,                              /* tp_getattr */

	0,                              /* tp_setattr */

	0,                              /* tp_compare */

	0,                              /* tp_repr */

	0,                              /* tp_as_number */

	0,                              /* tp_as_sequence */

	0,                              /* tp_as_mapping */

	0,                              /* tp_hash */

	0,                              /* tp_call */

	0,                              /* tp_str */

	0,                              /* tp_getattro */

	0,                              /* tp_setattro */

	0,                              /* tp_as_buffer */

	Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */

	ZstdDecompressionWriter__doc,   /* tp_doc */

	0,                              /* tp_traverse */

	0,                              /* tp_clear */

	0,                              /* tp_richcompare */

	0,                              /* tp_weaklistoffset */

	0,                              /* tp_iter */

	0,                              /* tp_iternext */

	ZstdDecompressionWriter_methods,/* tp_methods */

	ZstdDecompressionWriter_members,/* tp_members */

	0,                              /* tp_getset */

	0,                              /* tp_base */

	0,                              /* tp_dict */

	0,                              /* tp_descr_get */

	0,                              /* tp_descr_set */

	0,                              /* tp_dictoffset */

	0,                              /* tp_init */

	0,                              /* tp_alloc */

	PyType_GenericNew,              /* tp_new */

};

void decompressionwriter_module_init(PyObject* mod) {

	Py_SET_TYPE(&ZstdDecompressionWriterType, &PyType_Type);

	if (PyType_Ready(&ZstdDecompressionWriterType) < 0) {

		return;

	}

}