#include <stdlib.h>
#include <string.h> /* memmove() */
#include <iostream>

#include <SST/SST_Concurrency.h>

#include <ZUtil/ZAlloc.hpp>

SST_Mutex AllocLock = SST_Concurrency_CreateMutex();

#if ZALLOC_CHECK_ALLOC

static ZAllocInfo* AllocList = NULL;

int ZAllocInfo::CurrentAllocCount = 0;
int ZAllocInfo::CurrentCheckpoint = 0;

void ZAllocAddTrack(uintptr_t _address, size_t _size, const char *_file, int _line, int _checkpoint)
{	
	ZAllocInfo *info = (ZAllocInfo*)malloc(sizeof(ZAllocInfo));
	size_t len = strlen(_file);

	info->Address = _address;
	info->Size = _size;

	if(len > sizeof(info->File)-1)
		len = sizeof(info->File)-1;
		
	memcpy(info->File, _file, len);
	info->File[len] = '\0';

	info->Line = _line;
	info->Checkpoint = _checkpoint;

	info->Next = NULL;

	//Mutex may be NULL if doing pre-main() init 
	if (AllocLock)
		SST_Concurrency_LockMutex(AllocLock);

	info->Next = AllocList;
	AllocList = info;

	#if	ZALLOC_EXTRA_SPAMMY
	std::cerr << "Inc Alloc Count " << ZAllocInfo::CurrentAllocCount << " -> " << (ZAllocInfo::CurrentAllocCount+1) << std::endl;
	#endif

	ZAllocInfo::CurrentAllocCount++;

	//Mutex may be NULL if doing pre-main() init 
	if (AllocLock)
		SST_Concurrency_UnlockMutex(AllocLock);
}

void ZAllocRemoveTrack(uintptr_t _address)
{
	//Mutex may be NULL if doing pre-main() init 
	if (AllocLock)
		SST_Concurrency_LockMutex(AllocLock);

	ZAllocInfo *prev = NULL;
	ZAllocInfo *cur = AllocList;

	while (cur != NULL)
	{
		if (cur->Address == _address)
		{
			if (prev == NULL) //i.e. removing head
				AllocList = cur->Next;
			else
				prev->Next = cur->Next;	//Remove "this" from the chain. link parent's next pointer to this one's next pointer so nobody has a reference to "this".

			free(cur);

			#if	ZALLOC_EXTRA_SPAMMY
			std::cerr << "Dec Alloc Count " << ZAllocInfo::CurrentAllocCount << " -> " << (ZAllocInfo::CurrentAllocCount-1) << std::endl;
			#endif

			ZAllocInfo::CurrentAllocCount--;

			break;
		}

		prev = cur;
		cur = cur->Next;
	}

	//Mutex may be NULL if doing pre-main() init
	if (AllocLock)
		SST_Concurrency_UnlockMutex(AllocLock);
}

void ZAllocDumpUnfreed(ZMemState *_state)
{
	SST_Concurrency_LockMutex(AllocLock);

	size_t totalSize = 0;
	int checkpoint = _state == NULL ? 0 : _state->Checkpoint;

	std::cerr << "--- ZAlloc Unfreed Memory Dump ---" << std::endl << std::endl;
	std::cerr << ZAllocInfo::CurrentAllocCount << " Current Unfreed Allocations:" << std::endl;

	ZAllocInfo *cur = AllocList;

	while (cur != NULL)
	{
		if (cur->Checkpoint >= checkpoint)
		{		      
			std::cerr << cur->File << "\n\tLINE: " << cur->Line << "\n\tADDRESS " << (void*)cur->Address << " (" << cur->Size << " bytes) unfreed!" << std::endl;
			totalSize += cur->Size;
		}

		cur = cur->Next;		
	}

	std::cerr << "-----------------------------------------------------------\n";
	std::cerr << "Total Unfreed: " << totalSize << " bytes" << std::endl << std::endl;

	std::cerr << "---  End Unfreed Memory Dump   ---" << std::endl;

	SST_Concurrency_UnlockMutex(AllocLock);
}

//Debug Operator New
void * operator new(size_t _size, const char *_file, int _line, int _checkpoint)
throw()
{
	void *ptr = (void*)malloc(_size);

	ZAllocAddTrack((uintptr_t)ptr, _size, _file, _line, _checkpoint);

	return ptr;
}

void * operator new[](size_t _size, const char* _file, int _line, int _checkpoint)
throw()
{
	void *ptr = (void*)malloc(_size);

	ZAllocAddTrack((uintptr_t)ptr, _size, _file, _line, _checkpoint);

	return ptr;
}

//Debug Delete
void operator delete(void *_ptr, const char* _file, int _line, int _checkpoint)
{	
	(void)_file;
	(void)_line;
	(void)_checkpoint; 

	ZAllocRemoveTrack((uintptr_t)_ptr);

	free(_ptr);
}

void operator delete[](void *_ptr, const char* _file, int _line, int _checkpoint)
{
	(void)_file;
	(void)_line;
	(void)_checkpoint;

	ZAllocRemoveTrack((uintptr_t)_ptr);

	free(_ptr);
}
#ifndef __sun /* TODO: some error with Solaris Studio C++ compiler. */
void operator delete(void *_ptr)
{	
	ZAllocRemoveTrack((uintptr_t)_ptr);

	free(_ptr);
}

void operator delete[](void *_ptr)
{	
	ZAllocRemoveTrack((uintptr_t)_ptr);

	free(_ptr);
}
#endif

#endif

/////////////////////////////////////
/* Static New Functions and Fields */
/////////////////////////////////////

//This contains all our statically allocated pointers
void **StaticAllocArray = NULL;

//This is how many we can hold
int StaticAllocCapacity = 0;

//This is how many we've got
int StaticAllocSize = 0;

//This iterates and deletes
void ZAllocDeleteStatic(void)
{
	int i;

	//Mutex may be NULL if doing pre-main() init 
	if (AllocLock)
		SST_Concurrency_LockMutex(AllocLock);

	//Iterate, delete, then delete our static alloc array
	for (i = 0; i < StaticAllocSize; i++)
		free(StaticAllocArray[i]);

	free(StaticAllocArray);

	StaticAllocArray = NULL;
	StaticAllocCapacity = 0;
	StaticAllocSize = 0;

	//Mutex may be NULL if doing pre-main() init
	if (AllocLock)
		SST_Concurrency_UnlockMutex(AllocLock);
}

void ZAllocAddStaticPtr(void *_ptr)
{
	//Mutex may be NULL if doing pre-main() init 
	if (AllocLock)
		SST_Concurrency_LockMutex(AllocLock);

	//If we need to, register with atexit to call ZAllocDeleteStatic
	if (StaticAllocArray == NULL)
		atexit(ZAllocDeleteStatic);

	//Check to see if StaticAllocArray is full
	if (StaticAllocSize >= StaticAllocCapacity )
	{
		//Allocate new one
		size_t newCapacity = StaticAllocCapacity == 0 ? 128 : StaticAllocCapacity * 2;
		void *temp = malloc(newCapacity * sizeof(void*));

		//Make sure we don't need to crash hard here
		if (temp == NULL)
		{
			std::cerr << "ZAlloc failed to allocate static tracking array!" << std::endl;
			exit(-1);
		}

		//Copy the old data in
		memmove(temp, StaticAllocArray, StaticAllocSize * sizeof(void*));

		//Free the old pointer
		free(StaticAllocArray);
		
		//Reset our pointer
		StaticAllocArray = (void**)temp;
	}

	//Add new void*
	StaticAllocArray[StaticAllocSize] = _ptr;

	//Increment size
	StaticAllocSize++;

	//Mutex may be NULL if doing pre-main() init
	if (AllocLock)
		SST_Concurrency_UnlockMutex(AllocLock);
}

//Static operator new
void * operator new(size_t _size, znew_static_t _staticAlloc)
{
	URFP(_staticAlloc);

	void *ptr = (void*)malloc(_size);

	ZAllocAddStaticPtr(ptr);

	return ptr;
}

void * operator new[](size_t _size, znew_static_t _staticAlloc)
{
	URFP(_staticAlloc);

	void *ptr = (void*)malloc(_size);

	ZAllocAddStaticPtr(ptr);

	return ptr;
}