/* Berkeley DB interface.

   Author: Michael McLay

   Hacked: Guido van Rossum

   Btree and Recno additions plus sequence methods: David Ely

   Hacked by Gustavo Niemeyer <niemeyer@conectiva.com> fixing recno

   support.

   XXX To do:

   - provide a way to access the various hash functions

   - support more open flags

   The windows port of the Berkeley DB code is hard to find on the web:

   www.nightmare.com/software.html

*/

#include "Python.h"

#ifdef WITH_THREAD

#include "pythread.h"

#endif

#include <sys/types.h>

#include <sys/stat.h>

#include <fcntl.h>

#ifdef HAVE_DB_185_H

#include <db_185.h>

#else

#include <db.h>

#endif

/* Please don't include internal header files of the Berkeley db package

   (it messes up the info required in the Setup file) */

typedef struct {

	PyObject_HEAD

	DB *di_bsddb;

	int di_size;	/* -1 means recompute */

	int di_type;

#ifdef WITH_THREAD

	PyThread_type_lock di_lock;

#endif

} bsddbobject;

static PyTypeObject Bsddbtype;

#define is_bsddbobject(v) ((v)->ob_type == &Bsddbtype)

#define check_bsddbobject_open(v, r) if ((v)->di_bsddb == NULL) \

               { PyErr_SetString(BsddbError, \

				 "BSDDB object has already been closed"); \

                 return r; }

static PyObject *BsddbError;

static PyObject *

newdbhashobject(char *file, int flags, int mode,

		int bsize, int ffactor, int nelem, int cachesize,

		int hash, int lorder)

{

	bsddbobject *dp;

	HASHINFO info;

	if ((dp = PyObject_New(bsddbobject, &Bsddbtype)) == NULL)

		return NULL;

	info.bsize = bsize;

	info.ffactor = ffactor;

	info.nelem = nelem;

	info.cachesize = cachesize;

	info.hash = NULL; /* XXX should derive from hash argument */

	info.lorder = lorder;

#ifdef O_BINARY

	flags |= O_BINARY;

#endif

	Py_BEGIN_ALLOW_THREADS

	dp->di_bsddb = dbopen(file, flags, mode, DB_HASH, &info);

	Py_END_ALLOW_THREADS

	if (dp->di_bsddb == NULL) {

		PyErr_SetFromErrno(BsddbError);

#ifdef WITH_THREAD

		dp->di_lock = NULL;

#endif

		Py_DECREF(dp);

		return NULL;

	}

	dp->di_size = -1;

	dp->di_type = DB_HASH;

#ifdef WITH_THREAD

	dp->di_lock = PyThread_allocate_lock();

	if (dp->di_lock == NULL) {

		PyErr_SetString(BsddbError, "can't allocate lock");

		Py_DECREF(dp);

		return NULL;

	}

#endif

	return (PyObject *)dp;

}

static PyObject *

newdbbtobject(char *file, int flags, int mode,

	      int btflags, int cachesize, int maxkeypage,

	      int minkeypage, int psize, int lorder)

{

	bsddbobject *dp;

	BTREEINFO info;

	if ((dp = PyObject_New(bsddbobject, &Bsddbtype)) == NULL)

		return NULL;

	info.flags = btflags;

	info.cachesize = cachesize;

	info.maxkeypage = maxkeypage;

	info.minkeypage = minkeypage;

	info.psize = psize;

	info.lorder = lorder;

	info.compare = 0; /* Use default comparison functions, for now..*/

	info.prefix = 0;

#ifdef O_BINARY

	flags |= O_BINARY;

#endif

	Py_BEGIN_ALLOW_THREADS

	dp->di_bsddb = dbopen(file, flags, mode, DB_BTREE, &info);

	Py_END_ALLOW_THREADS

	if (dp->di_bsddb == NULL) {

		PyErr_SetFromErrno(BsddbError);

#ifdef WITH_THREAD

		dp->di_lock = NULL;

#endif

		Py_DECREF(dp);

		return NULL;

	}

	dp->di_size = -1;

	dp->di_type = DB_BTREE;

#ifdef WITH_THREAD

	dp->di_lock = PyThread_allocate_lock();

	if (dp->di_lock == NULL) {

		PyErr_SetString(BsddbError, "can't allocate lock");

		Py_DECREF(dp);

		return NULL;

	}

#endif

	return (PyObject *)dp;

}

static PyObject *

newdbrnobject(char *file, int flags, int mode,

	      int rnflags, int cachesize, int psize, int lorder,

	      size_t reclen, u_char bval, char *bfname)

{

	bsddbobject *dp;

	RECNOINFO info;

	int fd;

	if ((dp = PyObject_New(bsddbobject, &Bsddbtype)) == NULL)

		return NULL;

	info.flags = rnflags;

	info.cachesize = cachesize;

	info.psize = psize;

	info.lorder = lorder;

	info.reclen = reclen;

	info.bval = bval;

	info.bfname = bfname;

#ifdef O_BINARY

	flags |= O_BINARY;

#endif

	/* This is a hack to avoid a dbopen() bug that happens when

	 * it fails. */

	fd = open(file, flags);

	if (fd == -1) {

		dp->di_bsddb = NULL;

	}

	else {

		close(fd);

		Py_BEGIN_ALLOW_THREADS

		dp->di_bsddb = dbopen(file, flags, mode, DB_RECNO, &info);

		Py_END_ALLOW_THREADS

	}

	if (dp->di_bsddb == NULL) {

		PyErr_SetFromErrno(BsddbError);

#ifdef WITH_THREAD

		dp->di_lock = NULL;

#endif

		Py_DECREF(dp);

		return NULL;

	}

	dp->di_size = -1;

	dp->di_type = DB_RECNO;

#ifdef WITH_THREAD

	dp->di_lock = PyThread_allocate_lock();

	if (dp->di_lock == NULL) {

		PyErr_SetString(BsddbError, "can't allocate lock");

		Py_DECREF(dp);

		return NULL;

	}

#endif

	return (PyObject *)dp;

}

static void

bsddb_dealloc(bsddbobject *dp)

{

#ifdef WITH_THREAD

	if (dp->di_lock) {

		PyThread_acquire_lock(dp->di_lock, 0);

		PyThread_release_lock(dp->di_lock);

		PyThread_free_lock(dp->di_lock);

		dp->di_lock = NULL;

	}

#endif

	if (dp->di_bsddb != NULL) {

		int status;

		Py_BEGIN_ALLOW_THREADS

		status = (dp->di_bsddb->close)(dp->di_bsddb);

		Py_END_ALLOW_THREADS

		if (status != 0)

			fprintf(stderr,

				"Python bsddb: close errno %d in dealloc\n",

				errno);

	}

	PyObject_Del(dp);

}

#ifdef WITH_THREAD

#define BSDDB_BGN_SAVE(_dp) \

	Py_BEGIN_ALLOW_THREADS PyThread_acquire_lock(_dp->di_lock,1);

#define BSDDB_END_SAVE(_dp) \

	PyThread_release_lock(_dp->di_lock); Py_END_ALLOW_THREADS

#else

#define BSDDB_BGN_SAVE(_dp) Py_BEGIN_ALLOW_THREADS 

#define BSDDB_END_SAVE(_dp) Py_END_ALLOW_THREADS

#endif

static Py_ssize_t

bsddb_length(bsddbobject *dp)

{

	check_bsddbobject_open(dp, -1);

	if (dp->di_size < 0) {

		DBT krec, drec;

		int status;

		int size = 0;

		BSDDB_BGN_SAVE(dp)

		for (status = (dp->di_bsddb->seq)(dp->di_bsddb,

						  &krec, &drec,R_FIRST);

		     status == 0;

		     status = (dp->di_bsddb->seq)(dp->di_bsddb,

						  &krec, &drec, R_NEXT))

			size++;

		BSDDB_END_SAVE(dp)

		if (status < 0) {

			PyErr_SetFromErrno(BsddbError);

			return -1;

		}

		dp->di_size = size;

	}

	return dp->di_size;

}

static PyObject *

bsddb_subscript(bsddbobject *dp, PyObject *key)

{

	int status;

	DBT krec, drec;

	char *data,buf[4096];

	int size;

	PyObject *result;

	recno_t recno;

	if (dp->di_type == DB_RECNO) {

		if (!PyArg_Parse(key, "i", &recno)) {

			PyErr_SetString(PyExc_TypeError,

					"key type must be integer");

			return NULL;

		}

		krec.data = &recno;

		krec.size = sizeof(recno);

	}

	else {

		if (!PyArg_Parse(key, "s#", &data, &size)) {

			PyErr_SetString(PyExc_TypeError,

					"key type must be string");

			return NULL;

		}

		krec.data = data;

		krec.size = size;

	}

        check_bsddbobject_open(dp, NULL);

	BSDDB_BGN_SAVE(dp)

	status = (dp->di_bsddb->get)(dp->di_bsddb, &krec, &drec, 0);

	if (status == 0) {

		if (drec.size > sizeof(buf)) data = malloc(drec.size);

		else data = buf;

		if (data!=NULL) memcpy(data,drec.data,drec.size);

	}

	BSDDB_END_SAVE(dp)

	if (data==NULL) return PyErr_NoMemory();

	if (status != 0) {

		if (status < 0)

			PyErr_SetFromErrno(BsddbError);

		else

			PyErr_SetObject(PyExc_KeyError, key);

		return NULL;

	}

	result = PyString_FromStringAndSize(data, (int)drec.size);

	if (data != buf) free(data);

	return result;

}

static int

bsddb_ass_sub(bsddbobject *dp, PyObject *key, PyObject *value)

{

	int status;

	DBT krec, drec;

	char *data;

	int size;

	recno_t recno;

	if (dp->di_type == DB_RECNO) {

		if (!PyArg_Parse(key, "i", &recno)) {

			PyErr_SetString(PyExc_TypeError,

					"bsddb key type must be integer");

			return -1;

		}

		krec.data = &recno;

		krec.size = sizeof(recno);

	}

	else {

		if (!PyArg_Parse(key, "s#", &data, &size)) {

			PyErr_SetString(PyExc_TypeError,

					"bsddb key type must be string");

			return -1;

		}

		krec.data = data;

		krec.size = size;

	}

	check_bsddbobject_open(dp, -1);

	dp->di_size = -1;

	if (value == NULL) {

		BSDDB_BGN_SAVE(dp)

		status = (dp->di_bsddb->del)(dp->di_bsddb, &krec, 0);

		BSDDB_END_SAVE(dp)

	}

	else {

		if (!PyArg_Parse(value, "s#", &data, &size)) {

			PyErr_SetString(PyExc_TypeError,

					"bsddb value type must be string");

			return -1;

		}

		drec.data = data;

		drec.size = size;

		BSDDB_BGN_SAVE(dp)

		status = (dp->di_bsddb->put)(dp->di_bsddb, &krec, &drec, 0);

		BSDDB_END_SAVE(dp)

	}

	if (status != 0) {

		if (status < 0)

			PyErr_SetFromErrno(BsddbError);

		else

			PyErr_SetObject(PyExc_KeyError, key);

		return -1;

	}

	return 0;

}

static PyMappingMethods bsddb_as_mapping = {

	(lenfunc)bsddb_length,		/*mp_length*/

	(binaryfunc)bsddb_subscript,	/*mp_subscript*/

	(objobjargproc)bsddb_ass_sub,	/*mp_ass_subscript*/

};

static PyObject *

bsddb_close(bsddbobject *dp)

{

	if (dp->di_bsddb != NULL) {

		int status;

		BSDDB_BGN_SAVE(dp)

		status = (dp->di_bsddb->close)(dp->di_bsddb);

		BSDDB_END_SAVE(dp)

		if (status != 0) {

			dp->di_bsddb = NULL;

			PyErr_SetFromErrno(BsddbError);

			return NULL;

		}

	}

	dp->di_bsddb = NULL;

	Py_INCREF(Py_None);

	return Py_None;

}

static PyObject *

bsddb_keys(bsddbobject *dp)

{

	PyObject *list, *item=NULL;

	DBT krec, drec;

	char *data=NULL,buf[4096];

	int status;

	int err;

	check_bsddbobject_open(dp, NULL);

	list = PyList_New(0);

	if (list == NULL)

		return NULL;

	BSDDB_BGN_SAVE(dp)

	status = (dp->di_bsddb->seq)(dp->di_bsddb, &krec, &drec, R_FIRST);

	if (status == 0) {

		if (krec.size > sizeof(buf)) data = malloc(krec.size);

		else data = buf;

		if (data != NULL) memcpy(data,krec.data,krec.size);

	}

	BSDDB_END_SAVE(dp)

	if (status == 0 && data==NULL) return PyErr_NoMemory();

	while (status == 0) {

		if (dp->di_type == DB_RECNO)

			item = PyInt_FromLong(*((int*)data));

		else

			item = PyString_FromStringAndSize(data,

							  (int)krec.size);

		if (data != buf) free(data);

		if (item == NULL) {

			Py_DECREF(list);

			return NULL;

		}

		err = PyList_Append(list, item);

		Py_DECREF(item);

		if (err != 0) {

			Py_DECREF(list);

			return NULL;

		}

		BSDDB_BGN_SAVE(dp)

		status = (dp->di_bsddb->seq)

			(dp->di_bsddb, &krec, &drec, R_NEXT);

		if (status == 0) {

			if (krec.size > sizeof(buf))

				data = malloc(krec.size);

			else data = buf;

			if (data != NULL)

				memcpy(data,krec.data,krec.size);

		}

		BSDDB_END_SAVE(dp)

		if (data == NULL) return PyErr_NoMemory();

	}

	if (status < 0) {

		PyErr_SetFromErrno(BsddbError);

		Py_DECREF(list);

		return NULL;

	}

	if (dp->di_size < 0)

		dp->di_size = PyList_Size(list); /* We just did the work */

	return list;

}

static PyObject *

bsddb_has_key(bsddbobject *dp, PyObject *args)

{

	DBT krec, drec;

	int status;

	char *data;

	int size;

	recno_t recno;

	if (dp->di_type == DB_RECNO) {

		if (!PyArg_ParseTuple(args, "i;key type must be integer",

				      &recno)) {

			return NULL;

		}

		krec.data = &recno;

		krec.size = sizeof(recno);

	}

	else {

		if (!PyArg_ParseTuple(args, "s#;key type must be string",

				      &data, &size)) {

			return NULL;

		}

		krec.data = data;

		krec.size = size;

	}

	check_bsddbobject_open(dp, NULL);

	BSDDB_BGN_SAVE(dp)

	status = (dp->di_bsddb->get)(dp->di_bsddb, &krec, &drec, 0);

	BSDDB_END_SAVE(dp)

	if (status < 0) {

		PyErr_SetFromErrno(BsddbError);

		return NULL;

	}

	return PyInt_FromLong(status == 0);

}

static PyObject *

bsddb_set_location(bsddbobject *dp, PyObject *key)

{

	int status;

	DBT krec, drec;

	char *data,buf[4096];

	int size;

	PyObject *result;

	recno_t recno;

	if (dp->di_type == DB_RECNO) {

		if (!PyArg_ParseTuple(key, "i;key type must be integer",

				      &recno)) {

			return NULL;

		}

		krec.data = &recno;

		krec.size = sizeof(recno);

	}

	else {

		if (!PyArg_ParseTuple(key, "s#;key type must be string",

				      &data, &size)) {

			return NULL;

		}

		krec.data = data;

		krec.size = size;

	}

	check_bsddbobject_open(dp, NULL);

	BSDDB_BGN_SAVE(dp)