/*

python-bz2 - python bz2 library interface

Copyright (c) 2002  Gustavo Niemeyer <niemeyer@conectiva.com>

Copyright (c) 2002  Python Software Foundation; All Rights Reserved

*/

#include "Python.h"

#include <stdio.h>

#include <bzlib.h>

#include "structmember.h"

#ifdef WITH_THREAD

#include "pythread.h"

#endif

static char __author__[] =

"The bz2 python module was written by:\n\

\n\

    Gustavo Niemeyer <niemeyer@conectiva.com>\n\

";

/* Our very own off_t-like type, 64-bit if possible */

/* copied from Objects/fileobject.c */

#if !defined(HAVE_LARGEFILE_SUPPORT)

typedef off_t Py_off_t;

#elif SIZEOF_OFF_T >= 8

typedef off_t Py_off_t;

#elif SIZEOF_FPOS_T >= 8

typedef fpos_t Py_off_t;

#else

#error "Large file support, but neither off_t nor fpos_t is large enough."

#endif

#define BUF(v) PyString_AS_STRING((PyStringObject *)v)

#define MODE_CLOSED   0

#define MODE_READ     1

#define MODE_READ_EOF 2

#define MODE_WRITE    3

#define BZ2FileObject_Check(v)	(Py_TYPE(v) == &BZ2File_Type)

#ifdef BZ_CONFIG_ERROR

#if SIZEOF_LONG >= 8

#define BZS_TOTAL_OUT(bzs) \

	(((long)bzs->total_out_hi32 << 32) + bzs->total_out_lo32)

#elif SIZEOF_LONG_LONG >= 8

#define BZS_TOTAL_OUT(bzs) \

	(((PY_LONG_LONG)bzs->total_out_hi32 << 32) + bzs->total_out_lo32)

#else

#define BZS_TOTAL_OUT(bzs) \

	bzs->total_out_lo32

#endif

#else /* ! BZ_CONFIG_ERROR */

#define BZ2_bzRead bzRead

#define BZ2_bzReadOpen bzReadOpen

#define BZ2_bzReadClose bzReadClose

#define BZ2_bzWrite bzWrite

#define BZ2_bzWriteOpen bzWriteOpen

#define BZ2_bzWriteClose bzWriteClose

#define BZ2_bzCompress bzCompress

#define BZ2_bzCompressInit bzCompressInit

#define BZ2_bzCompressEnd bzCompressEnd

#define BZ2_bzDecompress bzDecompress

#define BZ2_bzDecompressInit bzDecompressInit

#define BZ2_bzDecompressEnd bzDecompressEnd

#define BZS_TOTAL_OUT(bzs) bzs->total_out

#endif /* ! BZ_CONFIG_ERROR */

#ifdef WITH_THREAD

#define ACQUIRE_LOCK(obj) PyThread_acquire_lock(obj->lock, 1)

#define RELEASE_LOCK(obj) PyThread_release_lock(obj->lock)

#else

#define ACQUIRE_LOCK(obj)

#define RELEASE_LOCK(obj)

#endif

/* Bits in f_newlinetypes */

#define NEWLINE_UNKNOWN	0	/* No newline seen, yet */

#define NEWLINE_CR 1		/* \r newline seen */

#define NEWLINE_LF 2		/* \n newline seen */

#define NEWLINE_CRLF 4		/* \r\n newline seen */

/* ===================================================================== */

/* Structure definitions. */

typedef struct {

	PyObject_HEAD

	PyObject *file;

	char* f_buf;		/* Allocated readahead buffer */

	char* f_bufend;		/* Points after last occupied position */

	char* f_bufptr;		/* Current buffer position */

	int f_softspace;	/* Flag used by 'print' command */

	int f_univ_newline;	/* Handle any newline convention */

	int f_newlinetypes;	/* Types of newlines seen */

	int f_skipnextlf;	/* Skip next \n */

	BZFILE *fp;

	int mode;

	Py_off_t pos;

	Py_off_t size;

#ifdef WITH_THREAD

	PyThread_type_lock lock;

#endif

} BZ2FileObject;

typedef struct {

	PyObject_HEAD

	bz_stream bzs;

	int running;

#ifdef WITH_THREAD

	PyThread_type_lock lock;

#endif

} BZ2CompObject;

typedef struct {

	PyObject_HEAD

	bz_stream bzs;

	int running;

	PyObject *unused_data;

#ifdef WITH_THREAD

	PyThread_type_lock lock;

#endif

} BZ2DecompObject;

/* ===================================================================== */

/* Utility functions. */

static int

Util_CatchBZ2Error(int bzerror)

{

	int ret = 0;

	switch(bzerror) {

		case BZ_OK:

		case BZ_STREAM_END:

			break;

#ifdef BZ_CONFIG_ERROR

		case BZ_CONFIG_ERROR:

			PyErr_SetString(PyExc_SystemError,

					"the bz2 library was not compiled "

					"correctly");

			ret = 1;

			break;

#endif

		case BZ_PARAM_ERROR:

			PyErr_SetString(PyExc_ValueError,

					"the bz2 library has received wrong "

					"parameters");

			ret = 1;

			break;

		case BZ_MEM_ERROR:

			PyErr_NoMemory();

			ret = 1;

			break;

		case BZ_DATA_ERROR:

		case BZ_DATA_ERROR_MAGIC:

			PyErr_SetString(PyExc_IOError, "invalid data stream");

			ret = 1;

			break;

		case BZ_IO_ERROR:

			PyErr_SetString(PyExc_IOError, "unknown IO error");

			ret = 1;

			break;

		case BZ_UNEXPECTED_EOF:

			PyErr_SetString(PyExc_EOFError,

					"compressed file ended before the "

					"logical end-of-stream was detected");

			ret = 1;

			break;

		case BZ_SEQUENCE_ERROR:

			PyErr_SetString(PyExc_RuntimeError,

					"wrong sequence of bz2 library "

					"commands used");

			ret = 1;

			break;

	}

	return ret;

}

#if BUFSIZ < 8192

#define SMALLCHUNK 8192

#else

#define SMALLCHUNK BUFSIZ

#endif

#if SIZEOF_INT < 4

#define BIGCHUNK  (512 * 32)

#else

#define BIGCHUNK  (512 * 1024)

#endif

/* This is a hacked version of Python's fileobject.c:new_buffersize(). */

static size_t

Util_NewBufferSize(size_t currentsize)

{

	if (currentsize > SMALLCHUNK) {

		/* Keep doubling until we reach BIGCHUNK;

		   then keep adding BIGCHUNK. */

		if (currentsize <= BIGCHUNK)

			return currentsize + currentsize;

		else

			return currentsize + BIGCHUNK;

	}

	return currentsize + SMALLCHUNK;

}

/* This is a hacked version of Python's fileobject.c:get_line(). */

static PyObject *

Util_GetLine(BZ2FileObject *f, int n)

{

	char c;

	char *buf, *end;

	size_t total_v_size;	/* total # of slots in buffer */

	size_t used_v_size;	/* # used slots in buffer */

	size_t increment;       /* amount to increment the buffer */

	PyObject *v;

	int bzerror;

	int bytes_read;

	int newlinetypes = f->f_newlinetypes;

	int skipnextlf = f->f_skipnextlf;

	int univ_newline = f->f_univ_newline;

	total_v_size = n > 0 ? n : 100;

	v = PyString_FromStringAndSize((char *)NULL, total_v_size);

	if (v == NULL)

		return NULL;

	buf = BUF(v);

	end = buf + total_v_size;

	for (;;) {

		Py_BEGIN_ALLOW_THREADS

		while (buf != end) {

			bytes_read = BZ2_bzRead(&bzerror, f->fp, &c, 1);

			f->pos++;

			if (bytes_read == 0) break;

			if (univ_newline) {

				if (skipnextlf) {

					skipnextlf = 0;

					if (c == '\n') {

						/* Seeing a \n here with skipnextlf true means we

						 * saw a \r before.

						 */

						newlinetypes |= NEWLINE_CRLF;

						if (bzerror != BZ_OK) break;

						bytes_read = BZ2_bzRead(&bzerror, f->fp, &c, 1);

						f->pos++;

						if (bytes_read == 0) break;

					} else {

						newlinetypes |= NEWLINE_CR;

					}

				}

				if (c == '\r') {

					skipnextlf = 1;

					c = '\n';

				} else if (c == '\n')

					newlinetypes |= NEWLINE_LF;

			}

			*buf++ = c;

			if (bzerror != BZ_OK || c == '\n') break;

		}

		if (univ_newline && bzerror == BZ_STREAM_END && skipnextlf)

			newlinetypes |= NEWLINE_CR;

		Py_END_ALLOW_THREADS

		f->f_newlinetypes = newlinetypes;

		f->f_skipnextlf = skipnextlf;

		if (bzerror == BZ_STREAM_END) {

			f->size = f->pos;

			f->mode = MODE_READ_EOF;

			break;

		} else if (bzerror != BZ_OK) {

			Util_CatchBZ2Error(bzerror);

			Py_DECREF(v);

			return NULL;

		}

		if (c == '\n')

			break;

		/* Must be because buf == end */

		if (n > 0)

			break;

		used_v_size = total_v_size;

		increment = total_v_size >> 2; /* mild exponential growth */

		total_v_size += increment;

		if (total_v_size > INT_MAX) {

			PyErr_SetString(PyExc_OverflowError,

			    "line is longer than a Python string can hold");

			Py_DECREF(v);

			return NULL;

		}

		if (_PyString_Resize(&v, total_v_size) < 0)

			return NULL;

		buf = BUF(v) + used_v_size;

		end = BUF(v) + total_v_size;

	}

	used_v_size = buf - BUF(v);

	if (used_v_size != total_v_size)

		_PyString_Resize(&v, used_v_size);

	return v;

}

/* This is a hacked version of Python's

 * fileobject.c:Py_UniversalNewlineFread(). */

size_t

Util_UnivNewlineRead(int *bzerror, BZFILE *stream,

		     char* buf, size_t n, BZ2FileObject *f)

{

	char *dst = buf;

	int newlinetypes, skipnextlf;

	assert(buf != NULL);

	assert(stream != NULL);

	if (!f->f_univ_newline)

		return BZ2_bzRead(bzerror, stream, buf, n);

	newlinetypes = f->f_newlinetypes;

	skipnextlf = f->f_skipnextlf;

	/* Invariant:  n is the number of bytes remaining to be filled

	 * in the buffer.

	 */

	while (n) {

		size_t nread;

		int shortread;

		char *src = dst;

		nread = BZ2_bzRead(bzerror, stream, dst, n);

		assert(nread <= n);

		n -= nread; /* assuming 1 byte out for each in; will adjust */

		shortread = n != 0;	/* true iff EOF or error */

		while (nread--) {

			char c = *src++;

			if (c == '\r') {

				/* Save as LF and set flag to skip next LF. */

				*dst++ = '\n';

				skipnextlf = 1;

			}

			else if (skipnextlf && c == '\n') {

				/* Skip LF, and remember we saw CR LF. */

				skipnextlf = 0;

				newlinetypes |= NEWLINE_CRLF;

				++n;

			}

			else {

				/* Normal char to be stored in buffer.  Also

				 * update the newlinetypes flag if either this

				 * is an LF or the previous char was a CR.

				 */

				if (c == '\n')

					newlinetypes |= NEWLINE_LF;

				else if (skipnextlf)

					newlinetypes |= NEWLINE_CR;

				*dst++ = c;

				skipnextlf = 0;

			}

		}

		if (shortread) {

			/* If this is EOF, update type flags. */

			if (skipnextlf && *bzerror == BZ_STREAM_END)

				newlinetypes |= NEWLINE_CR;

			break;

		}

	}

	f->f_newlinetypes = newlinetypes;

	f->f_skipnextlf = skipnextlf;

	return dst - buf;

}

/* This is a hacked version of Python's fileobject.c:drop_readahead(). */

static void

Util_DropReadAhead(BZ2FileObject *f)

{

	if (f->f_buf != NULL) {

		PyMem_Free(f->f_buf);

		f->f_buf = NULL;

	}

}

/* This is a hacked version of Python's fileobject.c:readahead(). */

static int

Util_ReadAhead(BZ2FileObject *f, int bufsize)

{

	int chunksize;

	int bzerror;

	if (f->f_buf != NULL) {

		if((f->f_bufend - f->f_bufptr) >= 1)

			return 0;

		else

			Util_DropReadAhead(f);

	}

	if (f->mode == MODE_READ_EOF) {

		f->f_bufptr = f->f_buf;

		f->f_bufend = f->f_buf;

		return 0;

	}

	if ((f->f_buf = PyMem_Malloc(bufsize)) == NULL) {

		PyErr_NoMemory();

		return -1;

	}

	Py_BEGIN_ALLOW_THREADS

	chunksize = Util_UnivNewlineRead(&bzerror, f->fp, f->f_buf,

					 bufsize, f);

	Py_END_ALLOW_THREADS

	f->pos += chunksize;

	if (bzerror == BZ_STREAM_END) {

		f->size = f->pos;

		f->mode = MODE_READ_EOF;

	} else if (bzerror != BZ_OK) {

		Util_CatchBZ2Error(bzerror);

		Util_DropReadAhead(f);

		return -1;

	}

	f->f_bufptr = f->f_buf;

	f->f_bufend = f->f_buf + chunksize;

	return 0;

}

/* This is a hacked version of Python's

 * fileobject.c:readahead_get_line_skip(). */

static PyStringObject *

Util_ReadAheadGetLineSkip(BZ2FileObject *f, int skip, int bufsize)

{

	PyStringObject* s;

	char *bufptr;

	char *buf;

	int len;

	if (f->f_buf == NULL)

		if (Util_ReadAhead(f, bufsize) < 0)

			return NULL;

	len = f->f_bufend - f->f_bufptr;

	if (len == 0)

		return (PyStringObject *)

			PyString_FromStringAndSize(NULL, skip);

	bufptr = memchr(f->f_bufptr, '\n', len);

	if (bufptr != NULL) {

		bufptr++;			/* Count the '\n' */

		len = bufptr - f->f_bufptr;

		s = (PyStringObject *)

			PyString_FromStringAndSize(NULL, skip+len);

		if (s == NULL)

			return NULL;

		memcpy(PyString_AS_STRING(s)+skip, f->f_bufptr, len);

		f->f_bufptr = bufptr;

		if (bufptr == f->f_bufend)

			Util_DropReadAhead(f);

	} else {

		bufptr = f->f_bufptr;

		buf = f->f_buf;

		f->f_buf = NULL; 	/* Force new readahead buffer */

                s = Util_ReadAheadGetLineSkip(f, skip+len,

					      bufsize + (bufsize>>2));

		if (s == NULL) {

		        PyMem_Free(buf);

			return NULL;

		}

		memcpy(PyString_AS_STRING(s)+skip, bufptr, len);

		PyMem_Free(buf);

	}

	return s;

}

/* ===================================================================== */

/* Methods of BZ2File. */

PyDoc_STRVAR(BZ2File_read__doc__,

"read([size]) -> string\n\

\n\

Read at most size uncompressed bytes, returned as a string. If the size\n\

argument is negative or omitted, read until EOF is reached.\n\

");

/* This is a hacked version of Python's fileobject.c:file_read(). */

static PyObject *

BZ2File_read(BZ2FileObject *self, PyObject *args)

{

	long bytesrequested = -1;

	size_t bytesread, buffersize, chunksize;

	int bzerror;

	PyObject *ret = NULL;

	if (!PyArg_ParseTuple(args, "|l:read", &bytesrequested))

		return NULL;

	ACQUIRE_LOCK(self);

	switch (self->mode) {

		case MODE_READ:

			break;

		case MODE_READ_EOF:

			ret = PyString_FromString("");

			goto cleanup;

		case MODE_CLOSED:

			PyErr_SetString(PyExc_ValueError,

					"I/O operation on closed file");

			goto cleanup;

		default:

			PyErr_SetString(PyExc_IOError,

					"file is not ready for reading");

			goto cleanup;

	}

	if (bytesrequested < 0)

		buffersize = Util_NewBufferSize((size_t)0);

	else

		buffersize = bytesrequested;