/**

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

 */

/* A Python C extension for Zstandard. */

#if defined(_WIN32)

#define WIN32_LEAN_AND_MEAN

#include <Windows.h>

#elif defined(__APPLE__) || defined(__OpenBSD__) || defined(__FreeBSD__) || defined(__NetBSD__) || defined(__DragonFly__)

#include <sys/types.h>

#include <sys/sysctl.h>

#endif

#include "python-zstandard.h"

PyObject *ZstdError;

PyDoc_STRVAR(estimate_decompression_context_size__doc__,

"estimate_decompression_context_size()\n"

"\n"

"Estimate the amount of memory allocated to a decompression context.\n"

);

static PyObject* estimate_decompression_context_size(PyObject* self) {

	return PyLong_FromSize_t(ZSTD_estimateDCtxSize());

}

PyDoc_STRVAR(frame_content_size__doc__,

"frame_content_size(data)\n"

"\n"

"Obtain the decompressed size of a frame."

);

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

	static char* kwlist[] = {

		"source",

		NULL

	};

	Py_buffer source;

	PyObject* result = NULL;

	unsigned long long size;

#if PY_MAJOR_VERSION >= 3

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

#else

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

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

	}

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

	if (size == ZSTD_CONTENTSIZE_ERROR) {

		PyErr_SetString(ZstdError, "error when determining content size");

	}

	else if (size == ZSTD_CONTENTSIZE_UNKNOWN) {

		result = PyLong_FromLong(-1);

	}

	else {

		result = PyLong_FromUnsignedLongLong(size);

	}

finally:

	PyBuffer_Release(&source);

	return result;

}

PyDoc_STRVAR(frame_header_size__doc__,

"frame_header_size(data)\n"

"\n"

"Obtain the size of a frame header.\n"

);

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

	static char* kwlist[] = {

		"source",

		NULL

	};

	Py_buffer source;

	PyObject* result = NULL;

	size_t zresult;

#if PY_MAJOR_VERSION >= 3

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

#else

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

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

	}

	zresult = ZSTD_frameHeaderSize(source.buf, source.len);

	if (ZSTD_isError(zresult)) {

		PyErr_Format(ZstdError, "could not determine frame header size: %s",

			ZSTD_getErrorName(zresult));

	}

	else {

		result = PyLong_FromSize_t(zresult);

	}

finally:

	PyBuffer_Release(&source);

	return result;

}

PyDoc_STRVAR(get_frame_parameters__doc__,

"get_frame_parameters(data)\n"

"\n"

"Obtains a ``FrameParameters`` instance by parsing data.\n");

PyDoc_STRVAR(train_dictionary__doc__,

"train_dictionary(dict_size, samples, k=None, d=None, steps=None,\n"

"                 threads=None,notifications=0, dict_id=0, level=0)\n"

"\n"

"Train a dictionary from sample data using the COVER algorithm.\n"

"\n"

"A compression dictionary of size ``dict_size`` will be created from the\n"

"iterable of ``samples``. The raw dictionary bytes will be returned.\n"

"\n"

"The COVER algorithm has 2 parameters: ``k`` and ``d``. These control the\n"

"*segment size* and *dmer size*. A reasonable range for ``k`` is\n"

"``[16, 2048+]``. A reasonable range for ``d`` is ``[6, 16]``.\n"

"``d`` must be less than or equal to ``k``.\n"

"\n"

"``steps`` can be specified to control the number of steps through potential\n"

"values of ``k`` and ``d`` to try. ``k`` and ``d`` will only be varied if\n"

"those arguments are not defined. i.e. if ``d`` is ``8``, then only ``k``\n"

"will be varied in this mode.\n"

"\n"

"``threads`` can specify how many threads to use to test various ``k`` and\n"

"``d`` values. ``-1`` will use as many threads as available CPUs. By default,\n"

"a single thread is used.\n"

"\n"

"When ``k`` and ``d`` are not defined, default values are used and the\n"

"algorithm will perform multiple iterations - or steps - to try to find\n"

"ideal parameters. If both ``k`` and ``d`` are specified, then those values\n"

"will be used. ``steps`` or ``threads`` triggers optimization mode to test\n"

"multiple ``k`` and ``d`` variations.\n"

);

static char zstd_doc[] = "Interface to zstandard";

static PyMethodDef zstd_methods[] = {

	{ "estimate_decompression_context_size", (PyCFunction)estimate_decompression_context_size,

	METH_NOARGS, estimate_decompression_context_size__doc__ },

	{ "frame_content_size", (PyCFunction)frame_content_size,

	METH_VARARGS | METH_KEYWORDS, frame_content_size__doc__ },

	{ "frame_header_size", (PyCFunction)frame_header_size,

	METH_VARARGS | METH_KEYWORDS, frame_header_size__doc__ },

	{ "get_frame_parameters", (PyCFunction)get_frame_parameters,

	METH_VARARGS | METH_KEYWORDS, get_frame_parameters__doc__ },

	{ "train_dictionary", (PyCFunction)train_dictionary,

	METH_VARARGS | METH_KEYWORDS, train_dictionary__doc__ },

	{ NULL, NULL }

};

void bufferutil_module_init(PyObject* mod);

void compressobj_module_init(PyObject* mod);

void compressor_module_init(PyObject* mod);

void compressionparams_module_init(PyObject* mod);

void constants_module_init(PyObject* mod);

void compressionchunker_module_init(PyObject* mod);

void compressiondict_module_init(PyObject* mod);

void compressionreader_module_init(PyObject* mod);

void compressionwriter_module_init(PyObject* mod);

void compressoriterator_module_init(PyObject* mod);

void decompressor_module_init(PyObject* mod);

void decompressobj_module_init(PyObject* mod);

void decompressionreader_module_init(PyObject *mod);

void decompressionwriter_module_init(PyObject* mod);

void decompressoriterator_module_init(PyObject* mod);

void frameparams_module_init(PyObject* mod);

void zstd_module_init(PyObject* m) {

	/* python-zstandard relies on unstable zstd C API features. This means

	   that changes in zstd may break expectations in python-zstandard.

	   python-zstandard is distributed with a copy of the zstd sources.

	   python-zstandard is only guaranteed to work with the bundled version

	   of zstd.

	   However, downstream redistributors or packagers may unbundle zstd

	   from python-zstandard. This can result in a mismatch between zstd

	   versions and API semantics. This essentially "voids the warranty"

	   of python-zstandard and may cause undefined behavior.

	   We detect this mismatch here and refuse to load the module if this

	   scenario is detected.

	*/

	if (ZSTD_VERSION_NUMBER != 10404 || ZSTD_versionNumber() != 10404) {

		PyErr_SetString(PyExc_ImportError, "zstd C API mismatch; Python bindings not compiled against expected zstd version");

		return;

	}

	bufferutil_module_init(m);

	compressionparams_module_init(m);

	compressiondict_module_init(m);

	compressobj_module_init(m);

	compressor_module_init(m);

	compressionchunker_module_init(m);

	compressionreader_module_init(m);

	compressionwriter_module_init(m);

	compressoriterator_module_init(m);

	constants_module_init(m);

	decompressor_module_init(m);

	decompressobj_module_init(m);

	decompressionreader_module_init(m);

	decompressionwriter_module_init(m);

	decompressoriterator_module_init(m);

	frameparams_module_init(m);

}

#if defined(__GNUC__) && (__GNUC__ >= 4)

#  define PYTHON_ZSTD_VISIBILITY __attribute__ ((visibility ("default")))

#else

#  define PYTHON_ZSTD_VISIBILITY

#endif

#if PY_MAJOR_VERSION >= 3

static struct PyModuleDef zstd_module = {

	PyModuleDef_HEAD_INIT,

	"zstd",

	zstd_doc,

	-1,

	zstd_methods

};

PYTHON_ZSTD_VISIBILITY PyMODINIT_FUNC PyInit_zstd(void) {

	PyObject *m = PyModule_Create(&zstd_module);

	if (m) {

		zstd_module_init(m);

		if (PyErr_Occurred()) {

			Py_DECREF(m);

			m = NULL;

		}

	}

	return m;

}

#else

PYTHON_ZSTD_VISIBILITY PyMODINIT_FUNC initzstd(void) {

	PyObject *m = Py_InitModule3("zstd", zstd_methods, zstd_doc);

	if (m) {

		zstd_module_init(m);

	}

}

#endif

/* Attempt to resolve the number of CPUs in the system. */

int cpu_count() {

	int count = 0;

#if defined(_WIN32)

	SYSTEM_INFO si;

	si.dwNumberOfProcessors = 0;

	GetSystemInfo(&si);

	count = si.dwNumberOfProcessors;

#elif defined(__APPLE__)

	int num;

	size_t size = sizeof(int);

	if (0 == sysctlbyname("hw.logicalcpu", &num, &size, NULL, 0)) {

		count = num;

	}

#elif defined(__linux__)

	count = sysconf(_SC_NPROCESSORS_ONLN);

#elif defined(__OpenBSD__) || defined(__FreeBSD__) || defined(__NetBSD__) || defined(__DragonFly__)

	int mib[2];

	size_t len = sizeof(count);

	mib[0] = CTL_HW;

	mib[1] = HW_NCPU;

	if (0 != sysctl(mib, 2, &count, &len, NULL, 0)) {

		count = 0;

	}

#elif defined(__hpux)

	count = mpctl(MPC_GETNUMSPUS, NULL, NULL);

#endif

	return count;

}

size_t roundpow2(size_t i) {

	i--;

	i |= i >> 1;

	i |= i >> 2;

	i |= i >> 4;

	i |= i >> 8;

	i |= i >> 16;

	i++;

	return i;

}

/* Safer version of _PyBytes_Resize().

 *

 * _PyBytes_Resize() only works if the refcount is 1. In some scenarios,

 * we can get an object with a refcount > 1, even if it was just created

 * with PyBytes_FromStringAndSize()! That's because (at least) CPython

 * pre-allocates PyBytes instances of size 1 for every possible byte value.

 *

 * If non-0 is returned, obj may or may not be NULL.

 */

int safe_pybytes_resize(PyObject** obj, Py_ssize_t size) {

	PyObject* tmp;

	if ((*obj)->ob_refcnt == 1) {

		return _PyBytes_Resize(obj, size);

	}

	tmp = PyBytes_FromStringAndSize(NULL, size);

	if (!tmp) {

		return -1;

	}

	memcpy(PyBytes_AS_STRING(tmp), PyBytes_AS_STRING(*obj),

		PyBytes_GET_SIZE(*obj));

	Py_DECREF(*obj);

	*obj = tmp;

	return 0;

}