/* Copyright (c) 2023 Yann Weber */
/*
   Redistribution and use in source and binary forms, with or without
   modification, are permitted provided that the following conditions
   are met:

   - Redistributions of source code must retain the above copyright
   notice, this list of conditions and the following disclaimer.

   - Redistributions in binary form must reproduce the above copyright
   notice, this list of conditions and the following disclaimer in the
   documentation and/or other materials provided with the distribution.

   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
   ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
   LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
   A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE FOUNDATION OR
   CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
   EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
   PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
   PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
   LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
   NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
   SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <errno.h>
#include <limits.h>
#include <poll.h>
#include <stdio.h>
#include <unistd.h>

#include <rnnoise.h>

#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include "structmember.h"

/**The number of samples processed by RNNoise.process_frame()
 * @note There is two samples per frame
 * @note Initialized by PyInit_rnnoise module initialization function
 */
static int frame_size;

/**@brief Represent a RNNoise python object */
typedef struct
{
	PyObject_VAR_HEAD;
	/** The RNNoise structure */
	DenoiseState *st;
	/** Read only attribute on rnnoise frame size
	 * @note The frame size is in sample (2 bytes per sample)
	 */
	int frame_size;

	/** When iterating use this iterator as source */
	PyObject *in_iter;
	/** Previous partially processed buffer */
	PyObject *prev_buf_obj;
	/** Pointer on the part of the buffer left to copy */
	char *prev_buf;
	/** Size in bytes of the buffer left to copy */
	Py_ssize_t left_in_buf;
} PyRNNoise_t;

/**@brief RNNoise.__new__ */
PyObject* pyrnnoise_new(PyTypeObject *subtype, PyObject *args, PyObject* kwds)
{
	PyObject *ret = PyType_GenericNew(subtype, args, kwds);
	if(PyErr_Occurred())
	{
		return NULL;
	}
	return ret;
}

/**@brief RNNoise.__init__
 * @todo implement model support in argument
 */
int pyrnnoise_init(PyObject *_self, PyObject *args, PyObject *kwds)
{
	PyRNNoise_t *self = (PyRNNoise_t*)_self;
	self->st = rnnoise_create(NULL); // TODO implement model instead of NULL
	self->frame_size = frame_size;
	self->in_iter = NULL;
	self->prev_buf_obj = NULL;
	return 0;
}

/** @brief RNNoise.__del__ */
void pyrnnoise_del(PyObject *_self)
{
	PyRNNoise_t *self = (PyRNNoise_t*)_self;
	rnnoise_destroy(self->st);
	if(self->prev_buf_obj) { Py_DECREF(self->prev_buf_obj); }
	if(self->in_iter) { Py_DECREF(self->in_iter); }
	return;
}

/**@brief RNNoise.process_frame(frame:bytes) -> bytes method
 * @param PyObject* Current RNNoise instance
 * @param PyObject** List of arguments (1 expected)
 * @param Py_ssize_t Number of arguments given
 * @return A PyObject* with processed frame (bytes) or NULL (on error)
 */
PyObject *pyrnnoise_process_frame(PyObject *_self, PyObject *const *args, Py_ssize_t nargs)
{
	PyRNNoise_t *self = (PyRNNoise_t*)_self;
	if(nargs != 1)
	{
		PyErr_SetString(PyExc_ValueError, "Expected one bytes argument");
		return NULL;
	}
	if(!PyBytes_Check(args[0]))
	{
		PyErr_SetString(PyExc_TypeError, "Excpected argument to be bytes");
		return NULL;
	}
	Py_ssize_t sz = PyBytes_Size(args[0]);
	if(sz != 2*self->frame_size)
	{
		PyErr_Format(PyExc_ValueError,
				"Expected %ld bytes but got %ld",
				self->frame_size*2, sz);
		return NULL;
	}

	short *data = (short*)PyBytes_AsString(args[0]);
	short out[self->frame_size];
	float frame[self->frame_size];

	for(int i=0; i<self->frame_size; i++) { frame[i] = data[i]; }
	rnnoise_process_frame(self->st, frame, frame);

	for(int i=0; i<self->frame_size; i++) { out[i] = frame[i]; }
	PyObject *ret = PyBytes_FromStringAndSize((char*)out, self->frame_size*2);
	return ret;
}

/**@brief pyrnnoise_iternext() utility function
 *
 * Fill the short buffer given in argument with previous buffer if set
 * or with next buffer(s) in iterator.
 *
 * @param PyRNNoise_t self
 * @param short* The buffer (with @ref frame_size samples)
 * @return The number of bytes read. Should be @ref frame_size * 2,
 * may be less than this value a last time before falling to 0 on
 * iterator exhauste
 */
static ssize_t _iter_frame(PyRNNoise_t *self, short *buf)
{
	const ssize_t to_read = frame_size * 2;

	PyObject *inbuf;
	ssize_t readed;

	if(self->prev_buf_obj)
	{
		Py_ssize_t from_buf = self->left_in_buf > to_read?\
				      to_read:self->left_in_buf;
		memcpy(buf, self->prev_buf, from_buf);
		readed = from_buf;
		if(from_buf == self->left_in_buf)
		{
			Py_DECREF(self->prev_buf_obj);
			self->prev_buf_obj = NULL;
			self->prev_buf = NULL;
		}
		else
		{
			self->left_in_buf -= from_buf;
			self->prev_buf = &self->prev_buf[from_buf];
		}
	}
	else
	{
		readed = 0;
	}

	while(readed < to_read)
	{
		const ssize_t left = to_read - readed;
		if(!(inbuf = PyIter_Next(self->in_iter)))
		{
			break;
		}
		if(!PyBytes_Check(inbuf))
		{
			PyErr_SetString(PyExc_TypeError, "Expected input iterator to return bytes");
			return -1;
		}
		Py_ssize_t inbuf_sz = PyBytes_Size(inbuf);
		char *inbuf_ptr = PyBytes_AsString(inbuf);

		memcpy((char*)buf+readed, inbuf_ptr,
				inbuf_sz > left?left:inbuf_sz);
		readed += inbuf_sz > left?left:inbuf_sz;
		if(inbuf_sz > left)
		{
			self->prev_buf = &inbuf_ptr[left];
			self->prev_buf_obj = inbuf;
			self->left_in_buf = inbuf_sz - left;
		}
		else
		{
			Py_DECREF(inbuf);
		}
	}

	return readed;
}

/**@brief The tp_iternext method of RNNoise
 * @return Next processed frame from input_iterator (bytes instance)
 */
PyObject *pyrnnoise_iternext(PyObject *_self)
{
	PyRNNoise_t *self = (PyRNNoise_t*)_self;
	if(!self->in_iter)
	{
		PyErr_SetString(PyExc_StopIteration, "No input use the iter_on class method");
		return NULL;
	}

	float frame[frame_size];
	short res[frame_size];

	ssize_t readed = _iter_frame(self, res);
	if(readed < 0)
	{
		return NULL;
	}
	else if(!readed)
	{
		Py_DECREF(self->in_iter);
		self->in_iter = NULL;
		PyErr_SetString(PyExc_StopIteration, "Input exhausted");
		return NULL;
	}
	else if(readed  % 2)
	{
		PyErr_SetString(PyExc_ValueError,
			"Iterator exhausted on invalid odd bytes count for as single channel 16bits PCM input");
		return NULL;
	}

	for(int i=0; i<readed/2; i++) { frame[i] = res[i]; }
	for(int i=readed/2; i<frame_size; i++) { frame[i] = 0; }

	rnnoise_process_frame(self->st, frame, frame);

	for(int i=0; i<frame_size; i++) { res[i] = frame[i]; }
	return PyBytes_FromStringAndSize((char*)res, readed);
}

/**@brief RNNoise.iter_on(iterator) method
 * @return A new reference on self iterator
 */
PyObject* pyrnnoise_iter_on(PyObject *_self, PyObject *iterator)
{
	PyRNNoise_t *self = (PyRNNoise_t*)_self;
	if(!PyIter_Check(iterator))
	{
		PyErr_SetString(PyExc_TypeError, "Expected an iterator");
		return NULL;
	}
	if(self->in_iter)
	{
		Py_DECREF(self->in_iter);
	}
	Py_INCREF(iterator);
	self->in_iter = iterator;
	return PyObject_GetIter(_self);
}


/*
 *	RNNoise class definition
 */

/**@brief RNNoise methods list */
PyMethodDef RNNoise_methods[] = {
	{"process_frame", (PyCFunction)pyrnnoise_process_frame, METH_FASTCALL,
		"Process a frame of data\n\
Arguments :\n\
    - frame : bytes representing a 16bit PCM audio signal (len must be \n\
              rnnoise.frame_size() * rnnoise.bytes_per_sample())\n\
Returns : The processed frame bytes"},
	{"iter_on", (PyCFunction)pyrnnoise_iter_on, METH_O,
		"Consume an iterator and iterate on processed frames\n\
Arguments :\n\
    - iterator : an iterator on PCM 16bit bytes\n\
Returns : an iterator on processed frames"},
	{ NULL } // Sentinel
};

/**@brief RNNoise attributes list */
PyMemberDef RNNoise_members[] = {
	{"frame_size", T_INT, offsetof(PyRNNoise_t, frame_size), READONLY,
	 "The number of sample processed by RNNoise.process_frame()"},
	{NULL} // Sentinel
};

/**@brief RNNoise type object */
PyTypeObject RNNoiseType = {
	PyVarObject_HEAD_INIT(NULL, 0)
	.tp_name = "rnnoise.RNNoise",
	.tp_doc= "Xiph RNNoise voice noise reduction wrapper",
	.tp_basicsize = sizeof(PyRNNoise_t),
	.tp_del = pyrnnoise_del,
	.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
	.tp_methods = RNNoise_methods,
	.tp_members = RNNoise_members,
	.tp_iter = PyObject_SelfIter,
	.tp_iternext = pyrnnoise_iternext,
	.tp_init = pyrnnoise_init,
	.tp_new = pyrnnoise_new,
};


/*
 *	rnnoise module definition
 */

/**@brief Module's method without argument, returning the number
 * of samples processed by RNNoise.process_frame()
 */
PyObject* pyrnnoise_frame_size(PyObject *mod, PyObject *nullval)
{
	return PyLong_FromLong(frame_size);
}


/**@brief Module's method without argument, returning the number of
 * bytes per sample
 */
PyObject* pyrnnoise_bytes_per_sample(PyObject *mod, PyObject *nullval)
{
	return PyLong_FromLong(2);
}

/**@brief rnnoise module's methods */
PyMethodDef pyrnnoise_methods[] = {
	{"frame_size", (PyCFunction)pyrnnoise_frame_size, METH_NOARGS,
		"Return the number of samples processed by RNNoise.process_frame() method"},
	{"bytes_per_sample", (PyCFunction)pyrnnoise_bytes_per_sample, METH_NOARGS,
		"Return the number of bytes per sample"},
	{NULL}
};

/**@brief rnnoise module object */
PyModuleDef pyrnnoise_module = {
	PyModuleDef_HEAD_INIT,
	.m_name = "rnnoise",
	.m_doc = "Python bindings to xiph rnnoise",
	.m_size = -1,
	.m_methods = pyrnnoise_methods,
	.m_slots = NULL,
	.m_traverse = NULL,
	.m_clear = NULL,
	.m_free = NULL,
};

/**@brief Module initialisation function (kind of main or __init__) */
PyMODINIT_FUNC
PyInit_rnnoise(void)
{
	PyObject *mod;

	mod = PyModule_Create(&pyrnnoise_module);
	if(mod == NULL) { return NULL; }

	if(PyType_Ready(&RNNoiseType) < 0)
	{
		goto fail_rnnoisetype_ready;
	}
	Py_INCREF(&RNNoiseType);
	if(PyModule_AddObject(mod, "RNNoise", (PyObject*)&RNNoiseType) < 0)
	{
		goto fail_rnnoisetype_add;
	}

	frame_size = rnnoise_get_frame_size();

	return mod;

fail_rnnoisetype_add:
	Py_DECREF(&RNNoiseType);
fail_rnnoisetype_ready:
	Py_DECREF(mod);
	return NULL;
}