#include "Python.h"
#include "compile.h"
#include "eval.h"
#include "frameobject.h"
#include "dictobject.h"

#ifdef MS_WINDOWS
#include <windows.h>
typedef __int64 hs_time;
#define GETTIMEOFDAY(P_HS_TIME) \
        { LARGE_INTEGER _temp; \
          QueryPerformanceCounter(&_temp); \
          *(P_HS_TIME) = _temp.QuadPart; }

typedef __int64 LONGLONG;
#else
#ifndef HAVE_GETTIMEOFDAY
#error "This module requires gettimeofday() on non-Windows platforms!"
#endif
#if (defined(PYOS_OS2) && defined(PYCC_GCC))
#include <sys/time.h>
#else
#include <sys/resource.h>
#include <sys/times.h>
#endif
#if defined(LINUX) || defined(SOLARIS)
typedef int64_t   LONGLONG;
#elif defined(HPUX)
#ifdef __STDC_32_MODE__
typedef long long          LONGLONG;
#else
typedef int64_t  LONGLONG;
#endif
#else
// add other cases as needed
typedef int64_t   LONGLONG;
#endif

typedef struct timeval hs_time;
#ifdef GETTIMEOFDAY_NO_TZ
#define GETTIMEOFDAY(ptv) gettimeofday((ptv))
#else
#define GETTIMEOFDAY(ptv) gettimeofday((ptv), (struct timezone *)NULL)
#endif
#endif

struct _ProfilerEntry;
/* Written by Brett Rosen and Ted Czotter */

/* The profiler could be made more efficient by doing the following:
 * 1 Replace current linked list structures and linear searches with
 *   a sorted array and binary search (O(n) -> O(ln(n))
 * 2 Embed the ProfilerEntry info into the PyCodeObject object eliminating
 *   the need for linear or binary search to locate the entry object.
 * 3 Change the interface of PyEval_SetProfile so that the callback
 *   can have the ProfilerContext allocated on the stack before calling
 *   the function and cleaned up after returning.
*/
typedef struct _ProfilerSubEntry // represents a function called from another function
{
    LONGLONG tt;
    LONGLONG it;
    long callcount;
    long recursivecallcount;
    long recursionLevel;
	struct _ProfilerEntry *entry;
	struct _ProfilerSubEntry *registryNext;
} ProfilerSubEntry;

typedef struct _ProfilerEntry // represents a function or user defined block
{
    LONGLONG tt; // total time in this entry
    LONGLONG it; // inline time in this entry (not in subcalls)
    long callcount; // how many times this was called
	long recursivecallcount; // how many times called recursively
    PyCodeObject *code; // Function object
    long recursionLevel;
	int nsubs;
    struct _ProfilerEntry *registryNext;
	ProfilerSubEntry* calls;
} ProfilerEntry;



typedef struct _ProfilerContext
{
    LONGLONG t0;
    LONGLONG subt;
    struct _ProfilerContext *previous;
	ProfilerEntry *ctxEntry;
} ProfilerContext;




typedef struct {
    PyObject_HEAD
	ProfilerEntry *profilerEntries;
	ProfilerContext *currentProfilerContext;
	LONGLONG hpFrequency;
	int enabled;
} ProfilerObject;

PyTypeObject PyProfiler_Type;

#define PyProfiler_Check(op) PyObject_TypeCheck(op, &PyProfiler_Type)
#define PyProfiler_CheckExact(op) ((op)->ob_type == &PyProfiler_Type)


static void addProfilerEntryToRegistry(ProfilerObject *pObj, ProfilerEntry *entry)
{
    entry->registryNext = pObj->profilerEntries;
	entry->calls = NULL;
	entry->nsubs = 0;
    pObj->profilerEntries = entry;
}

static void removeProfilerEntryFromRegistry(ProfilerObject *pObj, ProfilerEntry *entry)
{
    ProfilerEntry **ppe = &pObj->profilerEntries;
	ProfilerEntry *pe = NULL;
	for (pe = pObj->profilerEntries; pe; pe = pe->registryNext)
    {
		ProfilerSubEntry **ppse = &pe->calls;
		while (1)
		{
			ProfilerSubEntry *next = *ppse;
			if (next)
			{
				if (next->entry == entry)
				{
					*ppse = next->registryNext;
					free(next);
					pe->nsubs--;
					break;
				}
				ppse = &next->registryNext;
			}
			else
				break;
		}
    }
    while (1)
    {
        ProfilerEntry *next = *ppe;
        if (next)
        {
            if (next == entry)
            {
                *ppe = entry->registryNext;
                return;
            }
            ppe = &next->registryNext;
        }
        else
           return;
    }
}

static void initProfilerEntry(ProfilerObject *pObj, ProfilerEntry *self, PyCodeObject *_code)
{
	Py_INCREF(_code);
    self->code = _code;
    self->tt = 0;
    self->it = 0;
    self->callcount = 0;
	self->recursivecallcount = 0;
    self->recursionLevel = 0;
    addProfilerEntryToRegistry(pObj, self);
}
static void destroyProfilerEntry(ProfilerObject *pObj, ProfilerEntry *self)
{
	Py_DECREF(self->code);
    removeProfilerEntryFromRegistry(pObj, self);
}

static ProfilerEntry* findEntry(ProfilerObject *pObj, PyCodeObject *code)
{
	ProfilerEntry **ppe = &pObj->profilerEntries;
    while (1)
    {
        ProfilerEntry *next = *ppe;
        if (next)
        {
            if (next->code == code)
            {
				return next;
            }
            ppe = &next->registryNext;
        }
        else
		{
           next = (ProfilerEntry*) malloc(sizeof(ProfilerEntry));
		   initProfilerEntry(pObj, next, code);
		   return next;
		}
    }
}
static ProfilerSubEntry * getSubEntry(ProfilerObject *pObj,  ProfilerEntry *caller,
									  ProfilerEntry* entry, int create)
{
	ProfilerSubEntry **ppse = &caller->calls;
	while (1)
	{
		ProfilerSubEntry *next = *ppse;
		if (next)
		{
			if (next->entry == entry)
			{
				return next;
			}
			ppse = &next->registryNext;
		}
		else
		{
			if (!create)
				return NULL;

			next = (ProfilerSubEntry*) malloc(sizeof(ProfilerSubEntry));
			next->registryNext = caller->calls;
			next->entry = entry;
			next->tt = 0;
			next->it = 0;
			next->callcount = 0;
			next->recursivecallcount = 0;
			next->recursionLevel = 0;
			caller->calls = next;
			caller->nsubs++;
			return next;
		}
	}
}


static PyObject * profiler_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
	ProfilerObject *pRet = NULL;
#ifdef MS_WINDOWS
	LARGE_INTEGER li;
	int haveHPCounter;
#endif
	if (type != &PyProfiler_Type)
		return NULL;
	pRet = (ProfilerObject*)PyObject_New(ProfilerObject, type);
    if (pRet == NULL)
            return PyErr_NoMemory();

    pRet->profilerEntries = NULL;
	pRet->currentProfilerContext = NULL;
	pRet->enabled = 0;

#ifdef MS_WINDOWS
    haveHPCounter = QueryPerformanceFrequency(&li);
	if (haveHPCounter)
	{
		pRet->hpFrequency = li.QuadPart;
	}
#else
	pRet->hpFrequency = 1000;
#endif
    return (PyObject*) pRet;
}

static void clearEntries(ProfilerObject *pObj)
{
	ProfilerEntry *pe = pObj->profilerEntries;
    while (pe)
    {
        ProfilerEntry *next = pe->registryNext;
		ProfilerSubEntry *pse = pe->calls;
		while (pse)
		{
			ProfilerSubEntry *snext = pse->registryNext;
			free(pse);
			pse = snext;
		}
		pe->calls = NULL;
		pe->nsubs = 0;
		Py_DECREF(pe->code);
		free(pe);
		pe = next;
    }
}
static void profiler_dealloc(ProfilerObject *op)
{

	if (op->enabled)
		PyEval_SetProfile(NULL, NULL);
	clearEntries(op);
	PyObject_Del(op);
}



#ifdef MS_WINDOWS
static LONGLONG hpTimer()
{
    LARGE_INTEGER li;
    QueryPerformanceCounter(&li);
    return li.QuadPart;
}
#else
static LONGLONG hpTimer()
{
	struct timeval tv;
	LONGLONG ret;
    GETTIMEOFDAY(&tv);
	ret = tv.tv_sec * 1000000 + tv.tv_usec;
	return ret;
}
#endif


static void initContext(ProfilerObject *pObj, ProfilerContext *self, ProfilerEntry *entry)
{
	self->ctxEntry = entry;
    self->subt = 0;
    self->previous = pObj->currentProfilerContext;
	pObj->currentProfilerContext = self;
    ++entry->recursionLevel;
	if (pObj->enabled > 1 && self->previous)
	{
		// find or create an entry for me in my caller's entry
        ProfilerEntry *caller = self->previous->ctxEntry;
        ProfilerSubEntry *subentry = getSubEntry(pObj, caller, entry, 1);
		++subentry->recursionLevel;
	}
	self->t0 = hpTimer();

}
static void Stop(ProfilerObject *pObj, ProfilerContext *self, ProfilerEntry *entry)
{
    LONGLONG tt = hpTimer() - self->t0;
    LONGLONG it = tt - self->subt;
    if (self->previous)
        self->previous->subt += tt;
    pObj->currentProfilerContext = self->previous;
    if (--entry->recursionLevel == 0)
        entry->tt += tt;
	else
		++entry->recursivecallcount;
    entry->it += it;
    entry->callcount++;
	if (pObj->enabled > 1 && self->previous)
	{
		// find or create an entry for me in my caller's entry
        ProfilerEntry *caller = self->previous->ctxEntry;
        ProfilerSubEntry *subentry = getSubEntry(pObj, caller, entry, 0);
		if (subentry)
		{
			if (--subentry->recursionLevel == 0)
				subentry->tt += tt;
            else
				++subentry->recursivecallcount;
            subentry->it += it;
            ++subentry->callcount;

		}
	}
}

static int profiler_callback(PyObject *self, PyFrameObject *frame, int what,
                  PyObject *arg)
{
	ProfilerObject *pObj = (ProfilerObject*)self;
	ProfilerContext *pContext;
	ProfilerEntry *profEntry=NULL;
    switch (what) {
    case PyTrace_CALL:
	//case PyTrace_C_CALL:
		profEntry = findEntry(pObj, frame->f_code);
		pContext = (ProfilerContext*) malloc(sizeof(ProfilerContext));
		initContext(pObj, pContext, profEntry);
        break;
	case PyTrace_RETURN:
    //case PyTrace_C_RETURN:
	//case PyTrace_C_EXCEPTION:
		profEntry = findEntry(pObj, frame->f_code);
		pContext = pObj->currentProfilerContext;
		if (profEntry)
		{
				if (pContext)
					Stop(pObj, pContext, profEntry);
		}
		else
		{
			if (pContext)
				pObj->currentProfilerContext = pContext->previous;
		}

		if (pContext)
			free(pContext);
        break;
    default:
        /* should never get here */
        break;
    }
    return 0;
}




static PyObject * profiler_getstats(ProfilerObject *pObj)
{
	int num = 0;
	int len = 0;
	int i = 0;
	PyObject *pRet=NULL;
	ProfilerEntry *entry;
#ifdef MS_WINDOWS
	double factor = 1000.0 / pObj->hpFrequency;
#else
	double factor = 1.0 / pObj->hpFrequency;
#endif
    for (entry = pObj->profilerEntries; entry; entry = entry->registryNext)
    {
        if (entry->callcount)
            num++;
    }
	pRet = PyTuple_New(num);
	num = 0;
	for (entry = pObj->profilerEntries; entry; entry = entry->registryNext)
    {
        if (entry->callcount)
        {
			PyObject *info = PyDict_New();
			PyObject *label=NULL;
			PyCodeObject *f_code= entry->code;
			label = PyString_FromFormat("%s in %s at %d",
				PyString_AsString(f_code->co_name),
				PyString_AsString(f_code->co_filename), f_code->co_firstlineno);
			PyDict_SetItemString(info, "Label", label);
			PyDict_SetItemString(info, "CallCount", PyLong_FromLong(entry->callcount));
			PyDict_SetItemString(info, "RecursiveCallCount", PyLong_FromLong(entry->recursivecallcount));
			PyDict_SetItemString(info, "TotalTime", PyFloat_FromDouble(factor * entry->tt));
			PyDict_SetItemString(info, "InlineTime", PyFloat_FromDouble(factor * entry->it));

			len = entry->nsubs;
			if (len)
			{
				ProfilerSubEntry **ppse = &entry->calls;
				PyObject *pSub = PyTuple_New(len);

				i = 0;
				while (1)
				{
					ProfilerSubEntry *next = *ppse;
					if (next)
					{
						ProfilerSubEntry *sentry = next;
						PyObject *sinfo = PyDict_New();
						PyObject *slabel=NULL;
						PyCodeObject *f_code= sentry->entry->code;
						label = PyString_FromFormat("%s in %s at %d",
							PyString_AsString(f_code->co_name),
							PyString_AsString(f_code->co_filename), f_code->co_firstlineno);
						PyDict_SetItemString(sinfo, "Label", label);
						PyDict_SetItemString(sinfo, "CallCount", PyLong_FromLong(sentry->callcount));
						PyDict_SetItemString(sinfo, "RecursiveCallCount", PyLong_FromLong(sentry->recursivecallcount));
						PyDict_SetItemString(sinfo, "TotalTime", PyFloat_FromDouble(factor * sentry->tt));
						PyDict_SetItemString(sinfo, "InlineTime", PyFloat_FromDouble(factor * sentry->it));
						PyTuple_SetItem(pSub, i, sinfo);
						i++;
					}
					else
						break;
					ppse = &next->registryNext;

				}
				PyDict_SetItemString(info, "Calls", pSub);
			}

			PyTuple_SetItem(pRet, num, info);
			num++;


        }
    }
	return pRet;
}
static PyObject * profiler_enable(ProfilerObject *self, PyObject*args)
{
	int enable=0;
	if (!PyArg_ParseTuple(args, "i:enable", &enable))
		return NULL;
	PyEval_SetProfile((Py_tracefunc) profiler_callback, (PyObject*)self);
	if (enable == 1)
		self->enabled = 1;
	else
		self->enabled = 2;
	return PyInt_FromLong(1);
}

static void flush_unmatched(ProfilerObject *pObj)
{
	while (pObj->currentProfilerContext)
	{
		ProfilerContext *pContext = pObj->currentProfilerContext;
		ProfilerEntry *profEntry= pContext->ctxEntry;
		if (profEntry)
			Stop(pObj, pContext, profEntry);
		else 
			pObj->currentProfilerContext = pContext->previous;
		if (pContext)
			free(pContext);
	}

}
static PyObject * profiler_disable(ProfilerObject *self)
{
	PyEval_SetProfile(NULL, NULL);
	self->enabled = 0;
	flush_unmatched(self);
	return PyInt_FromLong(0);
}
static PyObject * profiler_clear(ProfilerObject *pObj)
{
    ProfilerEntry *entry;
	PyObject *values=NULL;

    for (entry = pObj->profilerEntries; entry; entry = entry->registryNext)
    {
		ProfilerSubEntry *pse = entry->calls;
        entry->callcount = 0;
		entry->recursivecallcount = 0;
        entry->tt = 0;
        entry->it = 0;
		while (pse)
		{
			ProfilerSubEntry *next = pse->registryNext;
			free(pse);
			pse = next;
		}
		entry->calls = NULL;
		entry->nsubs = 0;
    }

	return PyInt_FromLong(1);
}
static PyMethodDef
profiler_functions[] = {
    {"getstats",	(PyCFunction)profiler_getstats,   METH_NOARGS},
	{"enable",		(PyCFunction)profiler_enable,   METH_VARARGS},
	{"disable",		(PyCFunction)profiler_disable,   METH_NOARGS},
    {"clear",		(PyCFunction)profiler_clear,   METH_NOARGS},

    {NULL, NULL}
};

PyDoc_STRVAR(profiler_doc,
                         "implements profiler");

PyTypeObject PyProfiler_Type = {
        PyObject_HEAD_INIT(NULL)
        0,                                      /* ob_size */
        "Profiler",                           /* tp_name */
        sizeof(ProfilerObject) ,   /* tp_basicsize */
        0,                   /* tp_itemsize */
        (destructor)profiler_dealloc,         /* tp_dealloc */
        0,                                      /* tp_print */
        0,                                      /* tp_getattr */
        0,                                      /* tp_setattr */
        0,                                                          /* tp_compare */
        0,                              /* tp_repr */
        0,                                                      /* tp_as_number */
        0,                                      /* tp_as_sequence */
        0,                                      /* tp_as_mapping */
        0,                                                              /* tp_hash */
        0,                                      /* tp_call */
        0,                                                          /* tp_str */
        PyObject_GenericGetAttr,                /* tp_getattro */
        0,                                      /* tp_setattro */
        0,                                      /* tp_as_buffer */
        Py_TPFLAGS_DEFAULT | Py_TPFLAGS_CHECKTYPES |
                Py_TPFLAGS_BASETYPE,            /* tp_flags */
        profiler_doc,                         /* tp_doc */
        0,                                      /* tp_traverse */
        0,                                      /* tp_clear */
        0,                                      /* tp_richcompare */
        0,                                      /* tp_weaklistoffset */
        0,                                      /* tp_iter */
        0,                                      /* tp_iternext */
        profiler_functions,                     /* tp_methods */
        0,                                      /* tp_members */
        0,                                      /* tp_getset */
        0,                                      /* tp_base */
        0,                                      /* tp_dict */
        0,                                      /* tp_descr_get */
        0,                                      /* tp_descr_set */
        0,                                      /* tp_dictoffset */
        0,                                      /* tp_init */
        0,                                      /* tp_alloc */
        profiler_new,                         /* tp_new */
        PyObject_Del,                           /* tp_free */
};

static PyMethodDef moduleMethods[] =
{
        {NULL,                 NULL}
};

PyMODINIT_FUNC initlsprof(void)
{
	PyObject *module, *d;
	PyProfiler_Type.ob_type = &PyType_Type;
    module = Py_InitModule3("lsprof", moduleMethods, "Fast profiler");
	d = PyModule_GetDict(module);
	PyDict_SetItemString(d, "Profiler", (PyObject *)&PyProfiler_Type);


}