/*
 * Copyright (C) 2019 Weber Yann
 * 
 * This file is part of PyFCGI.
 * 
 * PyFCGI is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * any later version.
 * 
 * PyFCGI is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 * 
 * You should have received a copy of the GNU Affero General Public License
 * along with PyFCGI.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "python_pyfcgi.h"

/* Globals definitions */
/* libpyfcgi context */
libpyfcgi_context_t libpyfcgi = { NULL, NULL, NULL, NULL, 0, NULL, NULL, 0 };
/* Python module methods specs */
/**@todo Add doc in last field */
PyMethodDef pyfcgimodule_methods[] = {
	{"start_response", (PyCFunction)pyfcgi_start_response, METH_FASTCALL, NULL},
	{"write_body", (PyCFunction)pyfcgi_write_body, METH_FASTCALL, NULL},
	{NULL} // Sentinel
};
/* Python module specs */
PyModuleDef pyfcgimodule = {
	PyModuleDef_HEAD_INIT,
	"libpyfcgi",
	"Python librarie for PyFCGi",
	-1,
	pyfcgimodule_methods,
	NULL, NULL, NULL, NULL
};

/* Function definition */

/**@brief Forge & send default headers */
static void _default_headers()
{
	size_t sz;
	void *tmp;
	while(1)
	{
		sz = snprintf(libpyfcgi.heads_buf, libpyfcgi.heads_buf_sz,
			LIBPYFCGI_DEFAULT_HEADERS, libpyfcgi.status_buf);

		if(sz < libpyfcgi.heads_buf_sz) { break; }

		sz = ((sz>>11)+1)<<11; // 2048 alloc blocks
		tmp = realloc(libpyfcgi.heads_buf, sz);
		if(!tmp)
		{
			pyfcgi_log(LOG_ALERT,
				"Unable to realloc headers buffer : %s",
				strerror(errno));
			exit(PYFCGI_FATAL);
		}
		libpyfcgi.heads_buf = tmp;
		libpyfcgi.heads_buf_sz = sz;
	}
}

static void _set_status_buf()
{
	PyObject *bytes;

	short decref = 1;
	if(!libpyfcgi.status)
	{
		libpyfcgi.status = Py_BuildValue(LIBPYFCGI_DEFAULT_STATUS);
		return;
	}
	if(PyBytes_Check(libpyfcgi.status))
	{
		bytes = libpyfcgi.status;
		decref = 0;
	}
	else if(PyUnicode_Check(libpyfcgi.status))
	{
		
		bytes = PyUnicode_AsUTF8String(libpyfcgi.status);
		if(!bytes)
		{
			pyfcgi_log(LOG_ERR, "Unable to encode status using UTF8 codec");
			log_expt(LOG_ERR);
		}
	}
	else
	{
		bytes = PyObject_Bytes(libpyfcgi.status);
		if(!bytes)
		{
			pyfcgi_log(LOG_ERR, "Unable to convert status to bytes");
			log_expt(LOG_ERR);
		}
	}
	strncpy(libpyfcgi.status_buf, PyBytes_AsString(bytes), LIBPYFCGI_STATUS_SZ);
	libpyfcgi.status_buf[LIBPYFCGI_STATUS_SZ-1] = '\0';
	if(decref) { Py_DECREF(bytes); }
}

void libpyfcgi_set_headers_buf()
{
	PyObject *heads_iter, *head_tuple, *head_seq, **head, *bytes[2], *repr;
	Py_ssize_t tsz, n_head, sz, buf_sz, buf_left, buf_off;
	char errstr[256];
	char *head_str[2], *buf_ptr;
	int i;
	void *tmp;
	short content_type; // indicate if headers contains content-type field

	_set_status_buf();

	if(!libpyfcgi.headers)
	{
		_default_headers();
		return;
	}

	heads_iter = PyObject_GetIter(libpyfcgi.headers);
	if(!heads_iter)
	{
		pyfcgi_log(LOG_ERR, "Unable to get iterator from given headers");
		return;
	}
	Py_INCREF(heads_iter);

	n_head = 0;
	buf_ptr = libpyfcgi.heads_buf; //buffer pointer
	buf_sz = libpyfcgi.heads_buf_sz; //buffer size
	buf_left = buf_sz; //buffer size left
	buf_off = 0; //buffer pointer offset
	while((head_tuple = PyIter_Next(heads_iter)))
	{
		Py_INCREF(head_tuple);
		head_seq = PySequence_Fast(head_tuple, "Headers element is not a sequence");
		if((tsz = PySequence_Fast_GET_SIZE(head_seq)) != 2)
		{
			repr = PyObject_ASCII(head_tuple);
			Py_INCREF(repr);
			snprintf(errstr, 256,
				"Headers element expected to have len == 2, but got len %ld for element %ld : %s", tsz, n_head, PyUnicode_AsUTF8(repr));
			Py_DECREF(repr);
			PyErr_SetString(PyExc_TypeError, errstr);
			return;
		}
		head = PySequence_Fast_ITEMS(head_seq);
		for(i=0; i<2; i++)
		{
			if(PyUnicode_Check(libpyfcgi.status))
			{
				bytes[i] = PyUnicode_AsUTF8String(head[i]);
				if(!bytes[i])
				{
					//Error encoding header field
					pyfcgi_log(LOG_ERR, "Unable to encode header field using UTF8 codec");
					goto head_decode_err;
				}
			}
			else
			{
				bytes[i] = PyObject_Bytes(head[i]);
				if(!bytes[i])
				{
					pyfcgi_log(LOG_ERR, "Unable to convert header field to bytes");
					goto head_decode_err;
				}
			}
			Py_INCREF(bytes[i]);
			head_str[i] = PyBytes_AsString(bytes[i]);
			
		}
		if(!strcmp(head_str[0], "Content-Type"))
		{
			content_type = 1;
		}
		// Writing current headers field & value in header buffer
		while(1)
		{
			sz = snprintf(buf_ptr, buf_left,
				"%s: %s\r\n", head_str[0], head_str[1]);
			if(sz < buf_left)
			{
				buf_off += sz+1;
				buf_ptr += sz;
				break;
			}

			// 2048 alloc bloc
			sz = (((buf_sz+sz)>>11)+1)<<11;
			tmp = realloc(libpyfcgi.heads_buf, sz);
			if(!tmp)
			{
				pyfcgi_log(LOG_ALERT,
					"Unable to realloc headers buffer from iterator : %s",
					strerror(errno));
				exit(PYFCGI_FATAL);
			}
			buf_ptr = tmp + buf_off;
			buf_left = sz - buf_off;
			libpyfcgi.heads_buf = tmp;
			libpyfcgi.heads_buf_sz = sz;
		}

		n_head++;
		Py_DECREF(bytes[0]);
		Py_DECREF(bytes[1]);
		Py_DECREF(head_seq);
		Py_DECREF(head_tuple);
	}
	Py_DECREF(heads_iter);
	while(!content_type)
	{
		//Append text/html default content-type
		/**@todo Autodetect content type ?? **/
		sz = sizeof(LIBPYFCGI_DEFAULT_CTYPE)-1;
		if(sz < buf_left)
		{
			buf_off += sz +1;
			strncpy(buf_ptr, LIBPYFCGI_DEFAULT_CTYPE, buf_left);
			buf_ptr += sz;
			break;
		}
		sz = (((buf_sz+sz)>>11)+1)<<11;
		tmp = realloc(libpyfcgi.heads_buf, sz);
		if(!tmp)
		{
				pyfcgi_log(LOG_ALERT,
					"Unable to realloc headers buffer from default Content-Type",
					strerror(errno));
				exit(PYFCGI_FATAL);
		}
		buf_ptr = tmp+buf_off;
		buf_left = sz - buf_off;
		libpyfcgi.heads_buf = tmp;
		libpyfcgi.heads_buf_sz = sz;
	}
	
	return;
head_decode_err:
	Py_DECREF(head_seq);
	Py_DECREF(head_tuple);
	Py_DECREF(heads_iter);
}

void libpyfcgi_send_headers()
{
	size_t sz;
	int ret;

	if(!libpyfcgi.out)
	{
		// invalid context
		PyErr_SetString(PyExc_RuntimeError,
			"Trying to send headers but not in FCGI context");
		return;
	}

	if(libpyfcgi.headers_sent)
	{
		PyErr_SetString(PyExc_RuntimeError,
			"Headers allready sent...");
		return;
	}
	
	libpyfcgi_set_headers_buf();
	if(PyErr_Occurred())
	{
		/**@todo Better exception/error handling ? */
		log_expt(LOG_ALERT);
		exit(PYFCGI_FATAL);
	}
	sz = strlen(libpyfcgi.heads_buf);
	ret = FCGX_PutStr(libpyfcgi.heads_buf, sz, libpyfcgi.out);
	if(ret < 0 )
	{
		pyfcgi_log(LOG_ERR, "Unable to send headers");
		return;
	}
	libpyfcgi.rep_sz += ret;
	ret = FCGX_PutStr("\r\n", 2, libpyfcgi.out);
	if(ret < 0)
	{
		pyfcgi_log(LOG_ALERT, "Unable to send last \r\n from headers !");
	}
	libpyfcgi.rep_sz += ret;
	libpyfcgi.headers_sent = 1;
}

PyMODINIT_FUNC
PyInit_libpyfcgi(void)
{
	if(libpyfcgi.self != NULL)
	{
		return libpyfcgi.self;
	}

	// init IoIn type
	IoInType.tp_new = PyType_GenericNew;
	if(PyType_Ready(&IoInType) < 0)
	{
		return NULL;
	}

	// init module & globals
	libpyfcgi.status = NULL;
	libpyfcgi.headers = NULL;
	libpyfcgi.rep_sz = 0;
	libpyfcgi.self = PyModule_Create(&pyfcgimodule);
	if(libpyfcgi.self == NULL) { return NULL; }

	// Add type to module (optionnal)
	PyModule_AddObject(libpyfcgi.self, "IoIn", (PyObject*)&IoInType);

	// Create a new instance of IoIn
	libpyfcgi.ioin = (IoIn*)PyObject_CallObject((PyObject*)&IoInType, NULL);
	Py_INCREF(libpyfcgi.ioin);
	if(!libpyfcgi.ioin)
	{
		return NULL;
	}
	libpyfcgi.ioin->in_stream = &libpyfcgi.in; // point on stream pointer
	libpyfcgi.ioin->bin = 1; // binary stream (bytes for python)

	// Add it to wsgi dict
	if(PyDict_SetItemString(PyFCGI_conf.context.wsgi_dict, "wsgi.input",
		(PyObject*)libpyfcgi.ioin))
	{
		return NULL;
	}

	return libpyfcgi.self;
}

PyObject* pyfcgi_start_response(PyObject* self, PyObject** argv, Py_ssize_t argc)
{
	char err[128];
	PyObject *exc_info;
	if(argc != 2 && argc != 3)
	{
		snprintf(err, 128,
			"libpyfcgi.start_response() expected 2 to 3 arguments but %ld given",
			argc);
		PyErr_SetString(PyExc_ValueError, err);
		Py_RETURN_NONE;
	}

	if(libpyfcgi.status)
	{
		Py_DECREF(libpyfcgi.status);
	}
	if(libpyfcgi.headers)
	{
		Py_DECREF(libpyfcgi.headers);
	}

	libpyfcgi.status = argv[0];
	libpyfcgi.headers = argv[1];
	Py_INCREF(argv[0]);
	Py_INCREF(argv[1]);
	exc_info = argc>2?argv[2]:NULL;
	if(exc_info)
	{
		Py_INCREF(argv[2]);
		/* check if headers are sent, if yes re-raise the exception using
		   exc_info */
	}

	return PyObject_GetAttrString(self, "write_body");
}

PyObject* pyfcgi_write_body(PyObject* self, PyObject** argv, Py_ssize_t argc)
{
	char err[128];

	if(argc != 1)
	{
		snprintf(err, 128,
			"libpyfcgi.write_body() excpeted 1 argument but %ld given",
			argc);
		PyErr_SetString(PyExc_ValueError, err);
		Py_RETURN_NONE;
	}

	return _pyfcgi_write_body(argv[0]);
}

PyObject* _pyfcgi_write_body(PyObject *body_data)
{
	char err[128];
	const char *dat;
	PyObject *bytes, *cur, *iter, *repr;
	Py_ssize_t sz;
	int ret;


	if(!libpyfcgi.out)
	{
		// invalid context
		PyErr_SetString(PyExc_RuntimeError,
			"Trying to send body data but not in FCGI context");
		Py_RETURN_NONE;
	}

	cur = NULL;
	iter = NULL;
	if(PyUnicode_Check(body_data) || PyBytes_Check(body_data))
	{
		cur = body_data;
	}
	else if (!(iter = PyObject_GetIter(body_data)))
	{
		log_expt(LOG_ERR);
		repr = PyObject_ASCII(body_data);
		snprintf(err, 128, 
			"libpyfcgi.write_body() expect argument to be a str a bytes or an iterator, but %s given",
			PyUnicode_AsUTF8(repr));
		Py_DECREF(repr);
		PyErr_SetString(PyExc_ValueError, err);
		Py_RETURN_NONE;
	}
	else
	{
		Py_INCREF(iter);
		cur = PyIter_Next(iter);
	}

	if(!cur)
	{
		/**@todo trigger pyhon warning : empty body ? **/
		Py_RETURN_NONE;
	}

	// if headers not sent yet, send them....
	if(!libpyfcgi.headers_sent)
	{
		pyfcgi_log(LOG_DEBUG, "Headers not sent... sending them...");
		libpyfcgi_send_headers();
	}

	while(cur)
	{
		Py_INCREF(cur);
		bytes = NULL;
		if(PyUnicode_Check(cur))
		{
			dat = PyUnicode_AsUTF8AndSize(cur, &sz);
			if(!dat)
			{
				if(PyErr_Occurred()) { goto err_clean; }
				repr = PyObject_ASCII(cur);
				snprintf(err, 128,
					"libpyfcgi.__write_body unable to encode string as UTF-8 : %s",
					PyUnicode_AsUTF8(repr));
				Py_DECREF(repr);
				PyErr_SetString(PyExc_ValueError, err);
				goto err_clean;
			}
		}
		else if(PyBytes_Check(cur))
		{
			dat = PyBytes_AsString(cur);
			if(!dat)
			{
				if(PyErr_Occurred()) { goto err_clean; }

				snprintf(err, 128,
					"Unable to get bytes buffer when sending body data");
				PyErr_SetString(PyExc_ValueError, err);
				goto err_clean;
			}
			sz = PyBytes_GET_SIZE(cur);
		}
		else
		{
			bytes = PyObject_Bytes(cur);
			if(!bytes)
			{
				if(PyErr_Occurred()) { goto err_clean; }

				repr = PyObject_ASCII(cur);
				snprintf(err, 128,
					"libpyfcgi.__write_body expected str or bytes but %s given",
					PyUnicode_AsUTF8(repr));
				Py_DECREF(repr);
				PyErr_SetString(PyExc_ValueError, err);
				goto err_clean;
			}
			dat = PyBytes_AsString(cur);
			if(!dat)
			{
				if(PyErr_Occurred()) { goto err_clean; }

				snprintf(err, 128,
					"Unable to get bytes buffer when sending body data");
				PyErr_SetString(PyExc_ValueError, err);
				goto err_clean;
			}
			sz = PyBytes_GET_SIZE(cur);
		}
		
		ret = FCGX_PutStr(dat, sz, libpyfcgi.out);
		if(ret < 0)
		{
			pyfcgi_log(LOG_ALERT, "Unable to send reply");
			goto err_clean;
		}
		libpyfcgi.rep_sz += ret;

		if(bytes) { Py_DECREF(bytes); }
		Py_DECREF(cur);
		cur = iter?PyIter_Next(iter):NULL;
	}
	if(iter)
	{
		Py_DECREF(iter);
	}
	Py_RETURN_NONE;

err_clean:
	if(bytes) { Py_DECREF(bytes); }
	Py_DECREF(cur);
	if(iter) { Py_DECREF(iter); }
	Py_RETURN_NONE;
}