#include "munga.h"
#pragma hdrstop

#include "plug.h"
#include "node.h"
#include "hash.h"

class MemoryRegister;

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ MRStats ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

#define MR_FILENAME_LEN		(64)
#define MR_HASH_SIZE			(500)

class MRStats:
	public Plug
{
	friend class MemoryRegister;

public:
	MRStats(
		void *address,
      MRStatsAddressType address_type,
		char *fileName,
		int lineNum
	);
	~MRStats();

private:
	void
		*address;
	char
		fileName[MR_FILENAME_LEN];
	int
		lineNum;

   #ifdef MEMORY_VERIFY
      MRStatsAddressType
         addressType;
   #endif
};

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ MemoryRegister ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

class MemoryRegister:
	public Node
{
public:
	friend void Start_Registering_Function();
	friend void Stop_Registering_Function();

	friend void Register_Pointer_Function(
		void *address,
      MRStatsAddressType address_type,
		char *file_name,
		int line_num
	);
	friend void Unregister_Pointer_Function(
		void *address,
		char *file_name,
		int line_num
	);

	friend void Address_Registration_Function(
		void *address,
		char *file_name,
		int line_num
	);

	friend void Check_All_Signatures_Function(
		char *file_name,
		int line_num
	);

	MemoryRegister();
	~MemoryRegister();

	Logical
		TestInstance() const;

	void
		Report();

private:
	Logical
		RegisterAddress(
			void *address,
         MRStatsAddressType,
			char *fileName,
			int lineNum
		);

	Logical
		ReleaseAddress(
			void *address,
			char *fileName,
			int lineNum
		);

	Logical
		IsRegistered(
			void *address
		);

	void
		AddressRegistration(
			void *address,
			char *file_name,
			int line_num
		);

	void
		CheckAllSignatures(
			char *file_name,
			int line_num
		);

	static Logical
		VerifyAddress(
			void *address,
			char *fileName,
			int lineNum
		);

private:
	HashOf<MRStats*, unsigned long>
		memoryStats;
	Logical
		collectStats;
	long
		registrations;

	#ifdef MEMORY_VERIFY
		Logical
			checkAllSignatures;
		ChainOf<MRStats*>
			verificationSocket;
		Time
			nextTell;
	#endif
};

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~ MemoryRegisterShell ~~~~~~~~~~~~~~~~~~~~~~~~~~~~

class MemoryRegisterShell
{
public:
	MemoryRegisterShell();
   ~MemoryRegisterShell();

	static MemoryRegister
		*store;
};

#ifdef MEMORY_REGISTER
MemoryRegisterShell
	memory_register_shell;
#endif
MemoryRegister
	*MemoryRegisterShell::store=NULL;

//
//###########################################################################
// MemoryRegisterShell
//###########################################################################
//
MemoryRegisterShell::MemoryRegisterShell()
{
	store = new MemoryRegister;
	Check(store);
}

//
//###########################################################################
// ~MemoryRegisterShell
//###########################################################################
//
MemoryRegisterShell::~MemoryRegisterShell()
{
	MemoryRegister *obj;

	Check(store);
	store->Report();
	obj = store;
	store = NULL;
	delete obj;
}

//
//###########################################################################
// MRStats
//###########################################################################
//
#ifdef MEMORY_VERIFY
MRStats::MRStats(
	void *address,
   MRStatsAddressType address_type,
	char *fileName,
	int lineNum
)
{
	this->address = address;
	this->lineNum = lineNum;

	Str_Copy(this->fileName, fileName, MR_FILENAME_LEN);

   this->addressType = address_type;
}
#else
MRStats::MRStats(
	void *address,
   MRStatsAddressType,
	char *fileName,
	int lineNum
)
{
	this->address = address;
	this->lineNum = lineNum;

	Str_Copy(this->fileName, fileName, MR_FILENAME_LEN);
}
#endif

MRStats::~MRStats()
{
}

//
//###########################################################################
// MemoryRegister
//###########################################################################
//
#ifdef MEMORY_VERIFY
MemoryRegister::MemoryRegister():
	memoryStats(MR_HASH_SIZE, NULL, True),
   verificationSocket(NULL)
{
	collectStats = True;
	registrations = 0;
	checkAllSignatures = True;
	nextTell = Now();
}
#else
MemoryRegister::MemoryRegister():
	memoryStats(MR_HASH_SIZE, NULL, True)
{
	collectStats = True;
	registrations = 0;
}
#endif

//
//###########################################################################
// ~MemoryRegister
//###########################################################################
//
MemoryRegister::~MemoryRegister()
{
}

//
//###########################################################################
// TestInstance
//###########################################################################
//
Logical
	MemoryRegister::TestInstance() const
{
	Node::TestInstance();

	Check(&memoryStats);
	#ifdef MEMORY_VERIFY
		Check(&verificationSocket);
	#endif

	return True;
}

//
//###########################################################################
// RegisterAddress
//###########################################################################
//
Logical
	MemoryRegister::RegisterAddress(
		void *address,
		MRStatsAddressType address_type,
		char *fileName,
		int lineNum
	)
{
	MRStats *stats;

	if (!collectStats)
		return True;

	if (!VerifyAddress(address, fileName, lineNum))
		return False;

	Check(&memoryStats);
	if ((stats = memoryStats.Find((unsigned long)address)) != NULL)
	{
		Tell(
			stats->address << " : " <<
			stats->fileName << '(' <<
			stats->lineNum << ")\n"
		);
		Verify_Failed(
			"MemoryRegister::RegisterAddress: Address already registered by above",
			fileName,
			lineNum
		);
		return False;
	}

	stats = new MRStats(address, address_type, fileName, lineNum);
	Check(stats);

	Check(&memoryStats);
	collectStats = False;
	memoryStats.AddValue(stats, (unsigned long)address);
	#ifdef MEMORY_VERIFY
      // Currently, only verify signatured objects
		if (address_type == ObjectMRStatsAddressType)
		{
			Check(&verificationSocket);
			checkAllSignatures = False;
			verificationSocket.Add(stats);
			checkAllSignatures = True;
		}
	#endif
	collectStats = True;

	registrations++;

	return True;
}

//
//###########################################################################
// ReleaseAddress
//###########################################################################
//
Logical
	MemoryRegister::ReleaseAddress(
		void *address,
		char *fileName,
		int lineNum
	)
{
	MRStats *stats;

	if (!collectStats)
		return True;

	if (!VerifyAddress(address, fileName, lineNum))
		return False;

	stats = memoryStats.Find((unsigned long)address);
	if (stats != NULL)
	{
		Check(stats);
		collectStats = False;
		delete stats;
		collectStats = True;
		return True;
	}

	Verify_Failed(
		"MemoryRegister::ReleaseAddress: Address not found",
		fileName,
		lineNum
	);
	return False;
}

//
//###########################################################################
// IsRegistered
//###########################################################################
//
Logical
	MemoryRegister::IsRegistered(void *address)
{
	return (memoryStats.Find((unsigned long)address) != NULL);
}

//
//###########################################################################
// AddressRegistration
//###########################################################################
//
void 
	MemoryRegister::AddressRegistration(
		void *address,
		char *file_name,
		int line_num
	)
{
	MRStats *stats;

	if ((stats = memoryStats.Find((unsigned long)address)) != NULL)
	{
		std::cout <<
			stats->address << " : allocated at : " <<
			stats->fileName << '(' <<
         stats->lineNum << ")\n"
		;
	}
	else
	{
		Verify_Failed(
			"MemoryRegister::AddressRegistration: Address not found",
			file_name,
			line_num
		);
	}
}

//
//###########################################################################
// CheckAllSignatures
//###########################################################################
//
#ifdef MEMORY_VERIFY
void
	MemoryRegister::CheckAllSignatures(
		char *file_name,
		int line_num
   )
{
	if (checkAllSignatures)
	{
		checkAllSignatures = False;

		if (nextTell < Now())
		{
			nextTell += 1.0f;
			Tell(".");
		}

		ChainIteratorOf<MRStats*>
			iterator(&verificationSocket);
		MRStats
			*stats;

		Check(&iterator);
		while ((stats = iterator.ReadAndNext()) != NULL)
		{
			Check(stats);
			switch (stats->addressType)
			{
			case PointerMRStatsAddressType:
				Check_Pointer(stats->address);
				break;

			case ObjectMRStatsAddressType:
				if (Is_Signature_Bad((Signature*)stats->address))
				{
					Tell(
						stats->address << " : allocated at : " <<
						stats->fileName << '(' <<
						stats->lineNum << ") corrupted at : " <<
						file_name << '(' <<
						line_num << ")\n"
					);
					Check_Signature((Signature*)stats->address);
				}
				break;

			default:
				Fail("MemoryRegister::CheckAllSignatures - Unknown addressType");
				break;
			}
		}

		checkAllSignatures = True;
	}
}
#else
void
	MemoryRegister::CheckAllSignatures(
		char*,
		int
   )
{
}
#endif

//
//###########################################################################
// Report
//###########################################################################
//
void
	MemoryRegister::Report()
{
	HashIteratorOf<MRStats*, unsigned long>
   	iterator(memoryStats);
	MRStats
   	*stats;

	Tell("MemoryRegister::Report\n");
	Tell("Registrations " << registrations << std::endl);
	Tell("Unreleased Addresses:\n");

	while ((stats = iterator.ReadAndNext()) != NULL)
	{
		Check(stats);
		std::cout <<
			stats->address << " : " <<
         stats->fileName << '(' <<
         stats->lineNum << ")\n"
		;
	}
}

//
//###########################################################################
// VerifyAddress
//###########################################################################
//
Logical 
	MemoryRegister::VerifyAddress(
		void *address, 
		char *fileName, 
		int lineNum
	)
{
	//
	// HACK - An improvement would make sure that the memory
	// address is within our heap or stack 
	//
	if (address == 0L)
	{
		Verify_Failed(
			"MemoryRegister::VerifyAddress: Address is NULL", 
			fileName, 
			lineNum
		);
		return False;
	}
	return True;
}

//
//###########################################################################
// Start
//###########################################################################
//
void
	Start_Registering_Function()
{
	#if 0
	MemoryRegister::store = new MemoryRegister;
	Check(MemoryRegister::store);
   #endif
}

//
//###########################################################################
// Stop
//###########################################################################
//
void
	Stop_Registering_Function()
{
	#if 0
	MemoryRegister *obj;

	Check(MemoryRegister::store);
	MemoryRegister::store->Report();
	obj = MemoryRegister::store;
	MemoryRegister::store = NULL;
	delete obj;
   #endif
}

//
//###########################################################################
// Register_Pointer_Function
//###########################################################################
//
void
	Register_Pointer_Function(
		void *address,
      MRStatsAddressType address_type,
		char *file_name,
		int line_num
	)
{
	if (MemoryRegisterShell::store)
	{
		MemoryRegisterShell::store->RegisterAddress(
      	address,
         address_type,
         file_name,
         line_num
      );
	}
}

//
//###########################################################################
// Unregister_Pointer_Function
//###########################################################################
//
void
	Unregister_Pointer_Function(
		void *address,
		char *file_name,
		int line_num
	)
{
	if (MemoryRegisterShell::store)
	{
		MemoryRegisterShell::store->ReleaseAddress(address, file_name, line_num);
	}
}

//
//###########################################################################
// Address_Registration_Function
//###########################################################################
//
void
	Address_Registration_Function(
		void *address,
		char *file_name,
		int line_num
	)
{
	if (MemoryRegisterShell::store)
	{
		MemoryRegisterShell::store->AddressRegistration(address, file_name, line_num);
	}
}

//
//###########################################################################
// Check_All_Signatures_Function
//###########################################################################
//
void
	Check_All_Signatures_Function(
		char *file_name,
		int line_num
	)
{
	if (MemoryRegisterShell::store)
	{
		MemoryRegisterShell::store->CheckAllSignatures(file_name, line_num);
	}
}