#include "ZUnitTest.hpp"

#include <ZUtil/ZConcurrency.hpp>
#include <ZUtil/ZWeakPointer.hpp>

static const char* test_Constructors_Nullify();
static const char* test_Operators();
static const char* test_Cast();
static const char* test_ArrayContainment();
static const char* test_ListContainment();
static const char* test_RingBufferContainment();
static const char* test_DeallocLock();
static const char* test_MultithreadAccess();

//List of unit tests
ZUnitTest ZWeakPointerUnitTests[] =
{
	{ "ZWeakPointer: Constructors, Nullify",					test_Constructors_Nullify },
	{ "ZWeakPointer: Operators",								test_Operators },
	{ "ZWeakPointer: Cast",										test_Cast },
	{ "ZWeakPointer: Containment in ZArray",					test_ArrayContainment },
	{ "ZWeakPointer: Containment in ZList",						test_ListContainment },
	{ "ZWeakPointer: Containment in ZRingBuffer",				test_RingBufferContainment },
	{ "ZWeakPointer: Deallocation Lock",						test_DeallocLock },
	{ "ZWeakPointer: Multithreaded Access, Deallocation Lock",	test_MultithreadAccess }
};

DECLARE_ZTESTBLOCK(ZWeakPointer);

/*************************************************************************/

namespace ZWeakPointerTestObjects
{
	struct TestObject
	{
		//Constructor
		TestObject() { TestObject::Instances++; }

		//Destructor
		~TestObject() { TestObject::Instances--; }

		//Value field
		size_t Value;

		//Static Instance Count
		static size_t Instances;

		//Virtual Function
		virtual bool Foo() const { return true; }
	};

	struct TestObjectSubclass : public TestObject
	{
		//Constructor
		TestObjectSubclass() : TestObject() {}

		//Subclass Override
		bool Foo() const { return false; }

		//Subclass Method
		bool Bar() const { return true; }
	};

	//Defaults to zero
	size_t TestObject::Instances = 0;

	//This function is used to test the auto up-casting of contained types
	bool autoCast(ZWeakPointer<TestObject> _testObj)
	{
		return true;
	}

	//Method used to nullify a smart pointer in another thread
	int nullifySmartPointer(void* _sPtr)
	{
		ZSmartPointer<TestObject>& sPtr = *reinterpret_cast<ZSmartPointer<TestObject>*>(_sPtr);

		sPtr = NULL;

		return 0;
	}
}

using namespace ZWeakPointerTestObjects;

/*************************************************************************/

static const char* test_Constructors_Nullify()
{
	ZSmartPointer<TestObject> sNullPtr;
	ZSmartPointer<TestObject> sPtr(new TestObject());

	ZReferenceCounter* sPtrRefCounter = sPtr.Counter();

	ZWeakPointer<TestObject> wNullPtr;
	TASSERT(wNullPtr.Pointer() == NULL, "ZWeakPointer() constructed default pointer not set to null!\n");
	TASSERT(wNullPtr.Counter() == NULL, "ZWeakPointer() constructed default pointer with counter!\n");

	ZWeakPointer<TestObject> wNullPtrFromPtr(sNullPtr);
	TASSERT(wNullPtr.Pointer() == NULL, "ZWeakPointer() constructed from null smart pointer not set to null!\n");
	TASSERT(wNullPtr.Counter() == NULL, "ZWeakPointer() constructed from null smart pointer with counter!\n");

	ZWeakPointer<TestObject> wPtrFromSPtr = sPtr;
	TASSERT(TestObject::Instances == 1, "ZWeakPointer() constructed additional test object from smart pointer!");
	TASSERT(sPtr.Counter() == wPtrFromSPtr.Counter(), "ZWeakPointer() constructed from smart pointer does not share counter!");
	TASSERT(wPtrFromSPtr.Counter()->GetStrongRefCount() == 1, "ZWeakPointer() constructed from smart pointer added strong reference!");
	TASSERT(wPtrFromSPtr.Counter()->GetWeakRefCount() == 1, "ZWeakPointer() constructed from smart pointer did not add weak reference!");

	ZWeakPointer<TestObject> wPtrFromWPtr = wPtrFromSPtr;
	TASSERT(TestObject::Instances == 1, "ZWeakPointer() constructed additional test object from weak pointer!");
	TASSERT(wPtrFromSPtr.Counter() == wPtrFromWPtr.Counter(), "ZWeakPointer() constructed from weak pointer does not share counter!");
	TASSERT(wPtrFromWPtr.Counter()->GetStrongRefCount() == 1, "ZWeakPointer() constructed from weak pointer added strong reference!");
	TASSERT(wPtrFromWPtr.Counter()->GetWeakRefCount() == 2, "ZWeakPointer() constructed from weak pointer did not add weak reference!");

	//Equality Tests
	TASSERT(wPtrFromSPtr == sPtr, "Equality test failed on weak pointer from smart pointer and smart pointer!");
	TASSERT(wPtrFromWPtr == sPtr, "Equality test failed on weak pointer from weak pointer and smart pointer!");
	TASSERT(wPtrFromWPtr == wPtrFromSPtr, "Equality test failed on weak pointer from smart pointer and weak pointer from weak pointer!");

	//Lose the object
	sPtr = NULL;
	
	TASSERT(TestObject::Instances == 0, "Weak references prevented ZSmartPointer from deallocating object!");
	TASSERT(sPtr.Counter() == NULL, "ZSmartPointer failed to nullify pointer to reference counter!");

	TASSERT(sPtrRefCounter->GetStrongRefCount() == 0, "Nullification of smart pointer failed to lose strong reference!");

	TASSERT(sPtrRefCounter->GetWeakRefCount() == 2, "Nullification of smart pointer reduced weak reference count!");
	TASSERT(wPtrFromSPtr == NULL, "Nullification of weak pointer from smart pointer failed!");

	TASSERT(sPtrRefCounter->GetWeakRefCount() == 1, "Auto-nullification of weak pointer failed to lose weak reference!");
	TASSERT(wPtrFromWPtr == NULL, "Nullification of weak pointer from weak pointer failed!");

	TASSERT(wPtrFromSPtr.Counter() == NULL, "Nullification of weak pointer from smart pointer failed to lose reference counter!");
	TASSERT(wPtrFromWPtr.Counter() == NULL, "Nullification of weak pointer from weak pointer failed to lose reference counter!");

	return ZTEST_SUCCESS;
}

/*************************************************************************/

static const char* test_Operators()
{
	TestObject* object1 = new TestObject();
	TestObject* object2 = new TestObject();

	ZSmartPointer<TestObject> sPtr1(object1);
	ZSmartPointer<TestObject> sPtr2(object2);

	ZWeakPointer<TestObject> wPtr1 = sPtr1;
	ZWeakPointer<TestObject> wPtr2 = sPtr2;
	
	TASSERT(wPtr1 == object1, "Equality operator failed between weak pointer and object pointer!");
	TASSERT(wPtr1 == sPtr1, "Equality operator failed between weak pointer and smart pointer!");
	TASSERT(wPtr1 != sPtr2, "Equality operator succeeded between weak pointer and different smart pointer!");
	TASSERT(wPtr1 != wPtr2, "Equality operator succeeded between weak pointer and different weak pointer!");
	TASSERT(wPtr1->Foo(), "Access operator, call to Foo() failed to return true!");
	TASSERT((*wPtr1).Foo(), "Dereference operator, call to Foo() failed to return true!");

	return ZTEST_SUCCESS;
}

/*************************************************************************/

static const char* test_Cast()
{
	ZSmartPointer<TestObject> sPtr(new TestObjectSubclass());
	ZWeakPointer<TestObject> wPtr1(sPtr);
	ZWeakPointer<TestObjectSubclass> wPtr2(sPtr);

	TASSERT(TestObject::Instances == 1, "Subclass weak pointer created additional test object!");
	TASSERT(sPtr.Counter()->GetWeakRefCount() == 2, "Subclassed weak pointer failed to share reference counter!"); 

	TASSERT(!wPtr1->Foo(), "TestObject virtual method call failed!");
	TASSERT(!wPtr2->Foo(), "TestObjectSubclass virtual method call failed!");
	TASSERT(wPtr1.Cast<TestObjectSubclass>()->Bar(), "ZWeakPointer failed to cast!");

	TASSERT(autoCast(wPtr2), "ZWeakPointer failed to auto up-cast!");

	return ZTEST_SUCCESS;
}

/*************************************************************************/

static const char* test_ArrayContainment()
{
	ZArray< ZSmartPointer< TestObject > > sPtrArray(10);
	ZArray< ZWeakPointer< TestObject > > wPtrArray(10);

	TASSERT(TestObject::Instances == 0, "ZArray of smart pointers allocated instances!");

	for (size_t i = 0; i < 10; i++)
	{
		//Add a smart pointer to the set
		sPtrArray.PushBack(ZSmartPointer<TestObject>(znew TestObject()));
		sPtrArray[i]->Value = i;

		//Add a weak pointer created from that smart pointer
		wPtrArray.PushBack(sPtrArray[i]);
	}

	for (size_t i = 0; i < 10; i++)
	{
		TASSERT(sPtrArray[i].Counter()->GetStrongRefCount() == 1, "Starting Strong Count is incorrect!");
		TASSERT(sPtrArray[i].Counter()->GetWeakRefCount() == 1, "Starting Weak Count is incorrect!");
	}

	wPtrArray.PopBack();

	TASSERT(sPtrArray[sPtrArray.Size() - 1].Counter()->GetWeakRefCount() == 0, "PopBack() with weak pointer array failed to remove weak reference!");

	wPtrArray.PopFront();

	TASSERT(sPtrArray[0].Counter()->GetWeakRefCount() == 0, "PopFront() with weak pointer array failed to remove weak reference!");

	wPtrArray.Erase(0, 2);

	for (size_t i = 1; i < 3; i++)
	{
		TASSERT(sPtrArray[i].Counter()->GetWeakRefCount() == 0, "Erase() with weak pointer array failed to remove weak reference!");
	}

	//At this point, we have six elements in the weak pointer array
	wPtrArray.Resize(4);

	//That call should have ditched weak references from the last two
	for (size_t i = 7; i < 9; i++)
	{
		TASSERT(sPtrArray[i].Counter()->GetWeakRefCount() == 0, "Resize() with lower count failed to remove weak references!");
	}

	wPtrArray.Clear();

	for (size_t i = 0; i < 10; i++)
	{
		TASSERT(sPtrArray[i].Counter()->GetWeakRefCount() == 0, "Clear() failed to remove all weak references!");
	}

	return ZTEST_SUCCESS;
}

/*************************************************************************/

static const char* test_ListContainment()
{
	return "TODO";
}

/*************************************************************************/

static const char* test_RingBufferContainment()
{
	return "TODO";
}

/*************************************************************************/

static int useWeakPointer(void* _wPtr)
{
	ZWeakPointer<TestObject>& wPtr = *reinterpret_cast<ZWeakPointer<TestObject>*>(_wPtr);

	DeallocLock(wPtr) {
		return 0;
	} else {
		return 1;
	}
}

static const char* test_DeallocLock()
{
	TestObject* object = znew TestObject();

	ZSmartPointer<TestObject> sPtr(object);

	ZWeakPointer<TestObject> wPtr = sPtr;

	DeallocLock(wPtr) {
		TASSERT(wPtr.Counter()->State > 0, "ZWeakPointer DeallocLock did not increase State variable!");
	}

	TASSERT(wPtr.Counter()->State == 0, "ZWeakPointer DeallocLock did not decrement State variable!");

	ZThreadContext ctxt;
	ZThreadContext ctxt2;

	DeallocLock(wPtr) {
		ctxt = ZConcurrency::CreateThread(nullifySmartPointer, &sPtr);

		//At this point, the other thread should be waiting on this one, but 
		//should signal us by setting state to NULL with an atomic XOR
		while (wPtr.Counter()->State > 0)
			continue;

		ZWeakPointer<TestObject> wPtr2 = wPtr;

		ctxt2 = ZConcurrency::CreateThread(useWeakPointer, &wPtr2);

		ZConcurrency::WaitThread(ctxt2);

		//At the time we have gotten here, State has been signaled and we should fail to be able to signal 'InUse'
		TASSERT(ctxt2.ThreadExitStatus == 1, "ZWeakPointer was locked and used after being signaled for deallocation!");

		ZConcurrency::DestroyThread(ctxt2);
	}

	//Now that we have signaled that we are no longer in use, we can wait
	ZConcurrency::WaitThread(ctxt);
	ZConcurrency::DestroyThread(ctxt);

	TASSERT(TestObject::Instances == 0, "Failed to completely deallocate all TestObject instances!");

	return ZTEST_SUCCESS;
}

/*************************************************************************/

#define TEST_NUM_POINTERS (1024)

static ZArray< ZWeakPointer<TestObject> > g_threadPointers;

static int threadTest(void *arg)
{
	URFP(arg);

	while (WHILE_TRUE)
	{
		size_t numupdated = 0;

		for (size_t i = 0; i < TEST_NUM_POINTERS; i++)
		{
			if (g_threadPointers[i] == NULL)
				continue;

			//Effectively DeallocLock, but with the ability to inspect testLock
			if (ZWeakPointerDeallocLock testLock = g_threadPointers[i]()) {
				g_threadPointers[i]->Value++;
				numupdated++;
			}
		}

		if (numupdated == 0)
			break;
	}

	return 0;
}

static const char* test_MultithreadAccess()
{
	ZArray< ZSmartPointer<TestObject> > pointers;

	pointers.Resize(TEST_NUM_POINTERS);
	g_threadPointers.Resize(TEST_NUM_POINTERS);

	for (size_t i = 0; i < TEST_NUM_POINTERS; i++)
	{
		pointers[i] = znew TestObject();
		g_threadPointers[i] = pointers[i];

		pointers[i]->Value = 0;
	}

	ZThreadContext ctxt = ZConcurrency::CreateThread(threadTest, NULL);

	//Each object now has a smart pointer in this thread and a weak pointer that will be used 
	//by the other thread, so it should be safe to use the weak pointer as long as the DeallocLock
	//macro is used to prevent deallocation

	for (size_t i = 0; i < TEST_NUM_POINTERS; i++)
	{
		pointers[i] = NULL;

		ZConcurrency::SleepThread(10);
	}

	ZConcurrency::WaitThread(ctxt);
	ZConcurrency::DestroyThread(ctxt);

	TASSERT(TestObject::Instances == 0, "Failed to completely deallocate all TestObject instances!");

	return ZTEST_SUCCESS;
}