#include "RegressionTests.h"
#include "ThreadPool.h"
#include "PriorityQueueDLL.h"
#include <stdio.h>

static decltype(&InitializePriorityQueue) InitQueue;
static decltype(&EnqueueWithPriority) Enqueue;
static decltype(&DequeueHighest) Dequeue;
static decltype(&GetPriorityQueueCount) QueueCount;
static decltype(&IsPriorityQueueEmpty) QueueEmpty;
static decltype(&CleanupPriorityQueue) CleanupQueue;
static int Failures;

static void Check(bool condition, const char* name)
{
	if (!condition)
	{
		++Failures;
		printf("[FAIL] %s\n", name);
	}
}

static BOOL PoolEnqueue(PVOID queue, PNODE node, ULONG key)
{
	return Enqueue((PPRIORITY_QUEUE)queue, node, (PRIORITY)key);
}

static PNODE PoolDequeue(PVOID queue)
{
	return Dequeue((PPRIORITY_QUEUE)queue, NULL);
}

static void TestList()
{
	LIST_HEAD list;
	NODE a = {}, b = {}, absent = {};
	InitializeListHead(&list);
	Check(!RemoveHead(&list) && !RemoveTail(&list), "empty list removal");
	AddHead(&list, &a);
	AddTail(&list, &b);
	Check(list.pHead == &a && list.pTail == &b && a.pNext == &b && b.pPrev == &a, "list links");
	Check(SearchNode(&list, &a) == &a && !SearchNode(&list, &absent), "list search");
	Check(RemoveHead(&list) == &a && !a.pPrev && !a.pNext, "remove and detach head");
	Check(RemoveTail(&list) == &b && !list.pHead && !list.pTail, "remove last tail");
	AddHead(&list, &b);
	AddHead(&list, &a);
	Check(RemoveTail(&list) == &b && RemoveHead(&list) == &a, "opposite-end removal");
	DestroyListHead(&list);
}

static void TestQueue()
{
	PRIORITY_QUEUE queue;
	NODE nodes[6] = {};
	Check(!InitQueue(NULL), "null queue initialization");
	if (!InitQueue(&queue))
	{
		Check(false, "queue initialization");
		return;
	}
	Check(!Enqueue(&queue, NULL, PRIORITY_HIGH), "null node rejection");
	Check(!Enqueue(&queue, &nodes[0], (PRIORITY)-1), "negative priority rejection");
	Check(!Enqueue(&queue, &nodes[0], PRIORITY_COUNT), "invalid priority rejection");
	const PRIORITY priorities[] = { PRIORITY_LOW, PRIORITY_HIGH, PRIORITY_MEDIUM, PRIORITY_HIGH, PRIORITY_LOW, PRIORITY_MEDIUM };
	const int expected[] = { 1, 3, 2, 5, 0, 4 };
	for (int i = 0; i < 6; ++i)
	{
		Check(Enqueue(&queue, &nodes[i], priorities[i]) != FALSE, "queue insertion");
	}
	for (int i = 0; i < 6; ++i)
	{
		PRIORITY priority = PRIORITY_COUNT;
		Check(Dequeue(&queue, &priority) == &nodes[expected[i]], "highest priority and FIFO ordering");
		Check(priority == priorities[expected[i]], "dequeued priority");
		Check(QueueCount(&queue) == (ULONG)(5 - i), "queue count after dequeue");
	}
	Check(QueueEmpty(&queue) && !Dequeue(&queue, NULL), "drained queue");
	Check(Enqueue(&queue, &nodes[0], PRIORITY_HIGH) != FALSE, "reuse detached node");
	CleanupQueue(&queue);
	Check(!nodes[0].pNext && !nodes[0].pPrev, "cleanup detaches caller-owned nodes");
}

struct QueueRace
{
	PRIORITY_QUEUE Queue;
	NODE Nodes[256];
	volatile LONG ProducerDone;
	volatile LONG Consumed;
	volatile LONG Failed;
};

static DWORD WINAPI Produce(LPVOID context)
{
	QueueRace* race = (QueueRace*)context;
	for (int i = 0; i < 256; ++i)
	{
		if (!Enqueue(&race->Queue, &race->Nodes[i], (PRIORITY)(i % PRIORITY_COUNT)))
		{
			InterlockedExchange(&race->Failed, 1);
		}
	}
	InterlockedExchange(&race->ProducerDone, 1);
	return 0;
}

static DWORD WINAPI Consume(LPVOID context)
{
	QueueRace* race = (QueueRace*)context;
	while (!InterlockedCompareExchange(&race->ProducerDone, 0, 0) || !QueueEmpty(&race->Queue))
	{
		if (Dequeue(&race->Queue, NULL))
		{
			InterlockedIncrement(&race->Consumed);
		}
		if (QueueCount(&race->Queue) > 256)
		{
			InterlockedExchange(&race->Failed, 1);
		}
		SwitchToThread();
	}
	return 0;
}

static void TestConcurrentQueue()
{
	QueueRace race = {};
	if (!InitQueue(&race.Queue))
	{
		Check(false, "concurrent queue initialization");
		return;
	}
	HANDLE producer = CreateThread(NULL, 0, Produce, &race, 0, NULL);
	HANDLE consumer = CreateThread(NULL, 0, Consume, &race, 0, NULL);
	Check(producer && consumer, "queue test threads created");
	if (!producer) Produce(&race);
	if (!consumer) Consume(&race);
	if (producer) { WaitForSingleObject(producer, INFINITE); CloseHandle(producer); }
	if (consumer) { WaitForSingleObject(consumer, INFINITE); CloseHandle(consumer); }
	Check(race.Consumed == 256 && !race.Failed && QueueCount(&race.Queue) == 0, "concurrent queue count and delivery");
	CleanupQueue(&race.Queue);
}

struct Job
{
	HANDLE Started;
	HANDLE Release;
	volatile LONG Executions;
	volatile LONG Completions;
	volatile LONG Status;
};

static VOID Work(PVOID context)
{
	Job* job = (Job*)context;
	InterlockedIncrement(&job->Executions);
	if (job->Started) SetEvent(job->Started);
	if (job->Release) WaitForSingleObject(job->Release, INFINITE);
}

static VOID Complete(PVOID context, WORK_STATUS status)
{
	Job* job = (Job*)context;
	InterlockedExchange(&job->Status, status);
	InterlockedIncrement(&job->Completions);
}

struct ShutdownArgs
{
	THREADPOOL* Pool;
	HANDLE Returned;
};

static DWORD WINAPI StopPool(LPVOID context)
{
	ShutdownArgs* args = (ShutdownArgs*)context;
	ShutdownThreadPool(args->Pool);
	SetEvent(args->Returned);
	return 0;
}

static void TestShutdown(bool reap)
{
	PRIORITY_QUEUE queue;
	THREADPOOL pool;
	Job running = {};
	Job pending[20] = {};
	running.Started = CreateEvent(NULL, TRUE, FALSE, NULL);
	running.Release = CreateEvent(NULL, TRUE, FALSE, NULL);
	HANDLE returned = CreateEvent(NULL, TRUE, FALSE, NULL);
	if (!running.Started || !running.Release || !returned)
	{
		Check(false, "test events created");
		if (running.Started) CloseHandle(running.Started);
		if (running.Release) CloseHandle(running.Release);
		if (returned) CloseHandle(returned);
		return;
	}
	BOOL queueReady = InitQueue(&queue);
	BOOL poolReady = queueReady && InitializeThreadPool(&pool, 1, 1, &queue, PoolEnqueue, PoolDequeue, 30);
	Check(poolReady != FALSE, "test pool initialization");
	if (poolReady)
	{
		Check(!SubmitWork(&pool, NULL, NULL, 0, NULL), "null callback rejection");
		Check(!SubmitWork(&pool, Work, &running, PRIORITY_COUNT, Complete), "enqueue rejection propagation");
		Check(SubmitWork(&pool, Work, &running, 0, Complete) != FALSE, "running job accepted");
		Check(WaitForSingleObject(running.Started, 5000) == WAIT_OBJECT_0, "running job started");
		if (reap)
		{
			Sleep(50);
			Check(ReapStuckWorkers(&pool) == 1, "overdue job reported once");
			Check(ReapStuckWorkers(&pool) == 0 && running.Completions == 1, "repeated reap does not duplicate completion");
			Check(running.Status == WORK_STATUS_FAILED_STUCK, "stuck completion status");
		}
		else
		{
			for (Job& job : pending)
			{
				Check(SubmitWork(&pool, Work, &job, 0, Complete) != FALSE, "pending job accepted");
			}
			Check(GetPendingWorkCount(&pool) == 20, "blocked worker leaves jobs pending");
		}
		ShutdownArgs args = { &pool, returned };
		HANDLE shutdown = CreateThread(NULL, 0, StopPool, &args, 0, NULL);
		Check(shutdown != NULL, "shutdown thread created");
		if (shutdown)
		{
			bool stopping = false;
			ULONGLONG deadline = GetTickCount64() + 5000;
			while (!stopping && GetTickCount64() < deadline)
			{
				EnterCriticalSection(&pool.Lock);
				stopping = pool.Stop != 0;
				LeaveCriticalSection(&pool.Lock);
				if (!stopping) Sleep(1);
			}
			Check(stopping, "shutdown entered stopping state");
			Check(!SubmitWork(&pool, Work, &running, 0, Complete), "submission rejected during shutdown");
			Check(WaitForSingleObject(returned, 100) == WAIT_TIMEOUT, "shutdown waits for running callback");
			SetEvent(running.Release);
			WaitForSingleObject(shutdown, INFINITE);
			CloseHandle(shutdown);
		}
		else
		{
			SetEvent(running.Release);
			ShutdownThreadPool(&pool);
		}
		Check(running.Executions == 1 && running.Completions == 1, "running job executes and completes exactly once");
		Check(running.Status == (reap ? WORK_STATUS_FAILED_STUCK : WORK_STATUS_SUCCEEDED), "running job final status");
		Check(!pool.hWorkAvailable && !pool.hAllExited && pool.LiveThreads == 0, "shutdown joins workers and releases handles");
		if (!reap && shutdown)
		{
			for (Job& job : pending)
			{
				Check(job.Executions == 0 && job.Completions == 1 && job.Status == WORK_STATUS_CANCELLED, "pending job cancelled exactly once");
			}
		}
		Check(QueueEmpty(&queue), "shutdown empties work source");
	}
	if (queueReady) CleanupQueue(&queue);
	CloseHandle(running.Started);
	CloseHandle(running.Release);
	CloseHandle(returned);
}

int RunRegressionTests()
{
	Failures = 0;
	HMODULE dll = LoadLibraryExW(L"PriorityQueue.dll", NULL, LOAD_LIBRARY_SEARCH_APPLICATION_DIR | LOAD_LIBRARY_SEARCH_SYSTEM32);
	if (!dll)
	{
		printf("[FAIL] LoadLibrary PriorityQueue.dll: %lu\n", GetLastError());
		return 1;
	}
	InitQueue = (decltype(InitQueue))GetProcAddress(dll, "InitializePriorityQueue");
	Enqueue = (decltype(Enqueue))GetProcAddress(dll, "EnqueueWithPriority");
	Dequeue = (decltype(Dequeue))GetProcAddress(dll, "DequeueHighest");
	QueueCount = (decltype(QueueCount))GetProcAddress(dll, "GetPriorityQueueCount");
	QueueEmpty = (decltype(QueueEmpty))GetProcAddress(dll, "IsPriorityQueueEmpty");
	CleanupQueue = (decltype(CleanupQueue))GetProcAddress(dll, "CleanupPriorityQueue");
	if (!InitQueue || !Enqueue || !Dequeue || !QueueCount || !QueueEmpty || !CleanupQueue)
	{
		FreeLibrary(dll);
		printf("[FAIL] PriorityQueue exports\n");
		return 1;
	}
	TestList();
	TestQueue();
	TestConcurrentQueue();
	for (int i = 0; i < 10; ++i)
	{
		TestShutdown(false);
		TestShutdown(true);
	}
	FreeLibrary(dll);
	printf("REGRESSION TESTS: %s (%d failures)\n", Failures ? "FAILED" : "PASSED", Failures);
	return Failures ? 1 : 0;
}