/*
** Copyright 1999-2001 by LivingLogic AG, Bayreuth, Germany.
** Copyright 1999-2001 by Walter Dörwald
**
** All Rights Reserved
**
** Permission to use, copy, modify, and distribute this software and its documentation
** for any purpose and without fee is hereby granted, provided that the above copyright
** notice appears in all copies and that both that copyright notice and this permission
** notice appear in supporting documentation, and that the name of LivingLogic AG or
** the author not be used in advertising or publicity pertaining to distribution of the
** software without specific, written prior permission.
**
** LIVINGLOGIC AG AND THE AUTHOR DISCLAIM ALL WARRANTIES WITH REGARD TO THIS SOFTWARE,
** INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS, IN NO EVENT SHALL
** LIVING LOGIC AG OR THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL
** DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER
** IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR
** IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
*/

#include "Python.h"

#define ANSISTYLE_PRIMARY_COLORSTACK_SIZE 64

typedef struct
{
	PyObject_HEAD
	PyObject *stream;
	int stackSize;
	char activeColor;
	char primaryColorStack[ANSISTYLE_PRIMARY_COLORSTACK_SIZE];
} PyANSIStreamObject;

staticforward PyTypeObject PyANSIStream_Type;

#define PyANSIStream_Check(op) ((op)->ob_type == &PyANSIStream_Type)

extern int PyANSIStream_PushColor(PyANSIStreamObject *self, int color);
extern int PyANSIStream_PopColor(PyANSIStreamObject *self, int *result);

PyANSIStreamObject *PyANSIStream_New(PyObject *stream)
{
	PyANSIStreamObject *self;

	if (!stream)
		return NULL;

	if (!(self = PyObject_NEW(PyANSIStreamObject, &PyANSIStream_Type)))
		return NULL;

	Py_INCREF(stream);
	self->stream = stream;
	self->activeColor = 0x07;
	self->stackSize = 1;
	self->primaryColorStack[0] = 0x07;
	return (self);
}

static PyObject *ansistream_new(PyObject *self, PyObject *args)
{
	PyObject *stream;

	if (!PyArg_ParseTuple(args, "O", &stream))
		return NULL;

	return (PyObject *)PyANSIStream_New(stream);
}

static void ansistream_dealloc(PyANSIStreamObject *self)
{
	Py_XDECREF(self->stream);
	PyObject_DEL(self);
}

/* returns 0 on success */
static int _ansistream_getcolor(PyObject *object, int *color)
{
	PyObject *method;
	int colortemp;
	static PyObject *getcolorstring = NULL;
	PyObject *result;

	if (getcolorstring == NULL)
	{
		getcolorstring = PyString_InternFromString("getColor");
		if (getcolorstring == NULL)
			return (-1);
	}
	method = PyObject_GetAttr(object, getcolorstring);
	if (method == NULL)
		return (-1);

	result = PyEval_CallObject(method, NULL);
	Py_DECREF(method);
	if (!result)
		return (-1);
	if (!PyInt_Check(result))
	{
		PyErr_SetString(PyExc_TypeError, "getColor must return an integer");
		Py_DECREF(result);
		return (-1);
	}
	colortemp = PyInt_AsLong(result);
	Py_DECREF(result);
	if ((colortemp<-1 ) || (colortemp>0x7f))
	{
		PyErr_SetString(PyExc_ValueError, "color returned from getColor must be -1 or in range(0x7f)");
		return (-1);
	}
	if (color == NULL)
	{
		PyErr_BadInternalCall();
		return (0);
	}
	*color = colortemp;

	return (0);
}

/* returns 0 on success */
static int _ansistream_writestring(PyANSIStreamObject *self, PyStringObject *text)
{
	PyObject *method;
	static PyObject *writestring = NULL;
	PyObject *args;
	PyObject *result;

	if (writestring == NULL)
	{
		writestring = PyString_InternFromString("write");
		if (writestring == NULL)
			return (1);
	}
	if ((method = PyObject_GetAttr(self->stream, writestring)) == NULL)
		return (1);
	args = Py_BuildValue("(O)", text);
	if (!args)
	{
		Py_DECREF(method);
		return (1);
	}

	result = PyEval_CallObject(method, args);
	Py_DECREF(method);
	Py_DECREF(args);
	if (!result)
		return (1);
	Py_DECREF(result);

	return (0);
}

/* internal method: switches to the color color. If color is different from the currently active color (and not -1),
the appropriate ANSI escape sequence will be written to the stream.
returns 0 on success */
static int _ansistream_switchcolor(PyANSIStreamObject *self, int color)
{
	char activeColor = self->activeColor;

	if ((color != -1) && (color != activeColor))
	{
		char buffer[11];
		PyStringObject *string;
		char *p = buffer;
		int first = 1;

		*p++ = '\033';
		*p++ = '[';
		if ((!(color&0x08) && (activeColor&0x08)) /* do we have to get rid of the bold bit? (can only be done by a reset) */
			|| (color==0x07)) /* use reset when our target color is the default color (this is shorter than 40;37) */
		{
			*p++ = '0';
			activeColor = 0x07;
			first = 0;
		}
		/* now we know, that old and new color have the same boldness, or the new color is bold and the old isn't */
		if ((color&0x8) && !(activeColor&0x8))
		{
			if (!first)
				*p++ = ';';
			else
				first = 0;
			*p++ = '1';
			activeColor |= 0x08;
		}
		/* OK, now boldness is fixed, fix foreground now */
		if ((color&0x07) != (activeColor&0x07))
		{
			if (!first)
				*p++ = ';';
			else
				first = 0;
			*p++ = '3';
			*p++ = '0' + (color & 0x07);
			activeColor = (activeColor & 0x78) | (color & 0x07);
		}
		/* finally fix the background color */
		if ((color&0x70) != (activeColor&0x70))
		{
			if (!first)
				*p++ = ';';
			*p++ = '4';
			*p++ = '0' + ((color & 0x70) >> 4);
		}

		*p++ = 'm';

		string = (PyStringObject *)PyString_FromStringAndSize(buffer, p-buffer);
		if (string==NULL)
			return (-1);
		if (_ansistream_writestring(self, string))
		{
			Py_DECREF(string);
			return (-1);
		}
		Py_DECREF(string);
		self->activeColor = color;
	}
	return (0);
}

int PyANSIStream_Write(PyANSIStreamObject *self, PyObject *text)
{
	int i;

	if (PyString_Check(text) || PyUnicode_Check(text)) /* we assume that the stream knows how to handle unicode. */
	{
		if (_ansistream_switchcolor(self, self->primaryColorStack[self->stackSize-1]))
			return (-1);
		if (_ansistream_writestring(self, (PyStringObject *)text))
			return (-1);
	}
	else
	{
		int color;
		PyObject *content;
		static PyObject *contentstring = NULL;

		/* We could use the sequence method of the Text object directly, but going through __getitem__ for each access is to slow */
		if (contentstring == NULL)
		{
			contentstring = PyString_InternFromString("content");
			if (contentstring == NULL)
				return (-1);
		}
		content = PyObject_GetAttr(text, contentstring);
		if (content == NULL)
			return (-1);

		if (_ansistream_getcolor(text, &color))
			return (-1);
		if (color != -1)
		{
			if (PyANSIStream_PushColor(self, color))
			{
				Py_DECREF(content);
				return (-1);
			}
		}

		for (i = 0;;++i)
		{
			PyObject *item = PySequence_GetItem(content, i);
			if (item == NULL)
			{
				PyErr_Clear();
				break;
			}
			if (PyANSIStream_Write(self, item))
			{
				Py_DECREF(item);
				Py_DECREF(content);
				return (-1);
			}
			Py_DECREF(item);
		}
		if (color != -1)
		{
			if (PyANSIStream_PopColor(self, NULL))
			{
				Py_DECREF(content);
				return (-1);
			}
		}
	}
	return (0);
}

static PyObject *ansistream_write(PyANSIStreamObject *self, PyObject *args)
{
	PyObject *text;

	if (!PyArg_ParseTuple(args, "O:write", &text))
		return (NULL);

	if (PyANSIStream_Write(self, text))
		return (NULL);
	Py_INCREF(Py_None);
	return Py_None;
}

int PyANSIStream_PushColor(PyANSIStreamObject *self, int color)
{
	if (self->stackSize==ANSISTYLE_PRIMARY_COLORSTACK_SIZE)
	{
		PyErr_SetString(PyExc_MemoryError, "color stack overflow");
		return (-1);
	}

	if ((color<0 )|| (color>0x7f))
	{
		PyErr_SetString(PyExc_ValueError, "color must be in range(0x7f)");
		return (-1);
	}
	self->primaryColorStack[self->stackSize++] = color;
	return (0);
}

static PyObject *ansistream_pushcolor(PyANSIStreamObject *self, PyObject *args)
{
	int color;

	if (!PyArg_ParseTuple(args, "i:pushColor", &color))
		return (NULL);

	if (PyANSIStream_PushColor(self, color))
		return (NULL);

	Py_INCREF(Py_None);
	return (Py_None);
}

int PyANSIStream_PopColor(PyANSIStreamObject *self, int *result)
{
	if (self->stackSize==1)
	{
		PyErr_SetString(PyExc_ValueError, "color stack underflow");
		return (-1);
	}

	--self->stackSize;
	if (result)
		*result = self->primaryColorStack[self->stackSize];
	return (0);
}

static PyObject *ansistream_popcolor(PyANSIStreamObject *self, PyObject *args)
{
	int color;

	if (!PyArg_ParseTuple(args, ":popColor"))
		return (NULL);
	if (PyANSIStream_PopColor(self, &color))
		return (NULL);
	return (PyInt_FromLong(color));
}

int PyANSIStream_Finish(PyANSIStreamObject *self)
{
	return (_ansistream_switchcolor(self, 0x07));
}

static PyObject *ansistream_finish(PyANSIStreamObject *self, PyObject *args)
{
	if (!PyArg_ParseTuple(args, ":finish"))
		return (NULL);
	if (PyANSIStream_Finish(self))
		return (NULL);
	Py_INCREF(Py_None);
	return Py_None;
}

static PyObject *ansistream_repr(PyANSIStreamObject *self)
{
	char buf[300];
	sprintf(buf, "<Stream object at %lx>", (long)self);
	return PyString_FromString(buf);
}

static char ansistream_write__doc__[] =
"writes colored text to the underlying stream.";

static char ansistream_finish__doc__[] =
"this method should be called at the end of the output action to return to the default color.";

static char ansistream_pushcolor__doc__[] =
"pushes the passed in color (an integer between -1 and 0x7f) onto the color stack.";

static char ansistream_popcolor__doc__[] =
"removes the color from the top of the color stack and returns it.";
/* ==================================================================== */
/* ansistrem interface */
static PyMethodDef ansistream_methods[] =
{
	{"write", (PyCFunction)ansistream_write, METH_VARARGS, ansistream_write__doc__},
	{"finish", (PyCFunction)ansistream_finish, METH_VARARGS, ansistream_finish__doc__},
	{"pushColor", (PyCFunction)ansistream_pushcolor, METH_VARARGS, ansistream_pushcolor__doc__},
	{"popColor", (PyCFunction)ansistream_popcolor, METH_VARARGS, ansistream_popcolor__doc__},
	{NULL, NULL}
};

static PyObject *ansistream_getattr(PyANSIStreamObject *self, char *name)
{
	if (!strcmp(name, "stream"))
	{
		Py_INCREF(self->stream);
		return (self->stream);
	}
	else if (!strcmp(name, "__members__"))
		return Py_BuildValue("[s]", "stream");
	else
		return Py_FindMethod(ansistream_methods, (PyObject *)self, name);
}

static char stream__doc__[] =
"adds color capability to an output stream. A Stream\n\
keeps track of the current color and writes ANSI color escape\n\
sequences to the stream where appropriate. colors are numbers\n\
from 0 to 0x7f (or -1 for \"no color change\"). The high nybble\n\
of the color is the background color and the low nybble the\n\
foreground color.";

statichere PyTypeObject PyANSIStream_Type =
{
	PyObject_HEAD_INIT(NULL) /* will be fixed by the init function; required under Windows */
	0,
	"Stream",
	sizeof(PyANSIStreamObject),
	0,
	(destructor)ansistream_dealloc,   /* tp_dealloc */
	0,                                /* tp_print */
	(getattrfunc)ansistream_getattr,  /* tp_getattr */
	0,                                /* tp_setattr */
	0,                                /* tp_compare */
	(reprfunc)ansistream_repr,        /* tp_repr */
	0,                                /* tp_as_number */
	0,                                /* tp_as_sequence */
	0,                                /* tp_as_mapping */
	0,                                /* tp_hash */
	0,                                /* tp_call */
	0,                                /* tp_str */
	0,                                /* tp_getattro */
	0,                                /* tp_setattro */
	0,                                /* tp_asbuffer */
	Py_TPFLAGS_DEFAULT,               /* tp_flags */
	stream__doc__
};

/* ==================================================================== */
/* python module interface */

static PyMethodDef _functions[] =
{
	{"Stream", ansistream_new, METH_VARARGS, stream__doc__},
	{NULL, NULL}
};

static char module__doc__[] =
"This module contains the type Stream.";

void
#ifdef WIN32
__declspec(dllexport)
#endif
init_ansistyle(void)
{
	PyANSIStream_Type.ob_type = &PyType_Type; /* fix the type pointer */
	Py_InitModule3("_ansistyle", _functions, module__doc__);
}
