#include <Windows.h>
#include <stdio.h>
#include "Common.h"


// ==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==
// Global Variables

static PTP_POOL             g_Pool                  = NULL;
static TP_CALLBACK_ENVIRON  g_CallBackEnv           = { 0 };
static PTP_CLEANUP_GROUP    g_CleanupGroup          = NULL;
static volatile LONG64      g_l64ActiveWorkItems    = 0x00;      

// ==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==

BOOL InitializeThreadPool() 
{
    if (!(g_Pool = CreateThreadpool(NULL))) {
        printf("[!] CreateThreadpool Failed With Error: %lu\n", GetLastError());
        return FALSE;
    }

    if (!SetThreadpoolThreadMinimum(g_Pool, MIN_THREADS)) {
        printf("[!] SetThreadpoolThreadMinimum Failed With Error: %lu\n", GetLastError());
        return FALSE;
    }
    
    SetThreadpoolThreadMaximum(g_Pool, MAX_THREADS);
    InitializeThreadpoolEnvironment(&g_CallBackEnv);

    if (!(g_CleanupGroup = CreateThreadpoolCleanupGroup())) {
        printf("[!] CreateThreadpoolCleanupGroup Failed With Error: %lu\n", GetLastError());
        return FALSE;
    }
    
    SetThreadpoolCallbackPool(&g_CallBackEnv, g_Pool);
    SetThreadpoolCallbackCleanupGroup(&g_CallBackEnv, g_CleanupGroup, NULL);

    return TRUE;
}

// ==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==


BOOL DestroyThreadPool() 
{
    // Wait before cleaning up resources
    while (InterlockedCompareExchange64(&g_l64ActiveWorkItems, 0, 0) > 0)
    {
        Sleep(10);
    }

    if (g_CleanupGroup) 
    {
        CloseThreadpoolCleanupGroupMembers(g_CleanupGroup, FALSE, NULL);
        CloseThreadpoolCleanupGroup(g_CleanupGroup);
        g_CleanupGroup = NULL;
    }
    
    DestroyThreadpoolEnvironment(&g_CallBackEnv);

    if (g_Pool) {
        CloseThreadpool(g_Pool);
        g_Pool = NULL;
    }

	return TRUE;
}

// ==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==

VOID CALLBACK ThreadpoolCallback(PTP_CALLBACK_INSTANCE Instance, PVOID lpParameter, PTP_WORK Work) 
{
    UNREFERENCED_PARAMETER(Instance);

    PWORK_ITEM pWorkItem = (PWORK_ITEM)lpParameter;
    
    if (pWorkItem) 
    {
        pWorkItem->pRoutine(pWorkItem->pContext);
        HeapFree(GetProcessHeap(), 0x00, pWorkItem);
    }
    
    CloseThreadpoolWork(Work);
    InterlockedDecrement64(&g_l64ActiveWorkItems);
}

// ==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==-==

BOOL EnqueueThreadPoolWork(IN LPTHREAD_START_ROUTINE pRoutine, IN PVOID pContext)
{
    PTP_WORK    Work        = NULL;
    PWORK_ITEM  pWorkItem   = NULL;
    BOOL        bResult     = FALSE;

    if (!(pWorkItem = (PWORK_ITEM)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(WORK_ITEM)))) 
    {
		printf("[!] HeapAlloc Failed With Error: %lu\n", GetLastError());
        return FALSE;
    }
    
    pWorkItem->pRoutine = pRoutine;
    pWorkItem->pContext = pContext;

    if (!(Work = CreateThreadpoolWork(ThreadpoolCallback, pWorkItem, &g_CallBackEnv))) 
    {
        printf("[!] CreateThreadpoolWork Failed With Error: %lu\n", GetLastError());
        goto _END_OF_FUNC;
    }

    InterlockedIncrement64(&g_l64ActiveWorkItems);
    SubmitThreadpoolWork(Work);
    
    bResult = TRUE;

_END_OF_FUNC:
    if (!bResult && pWorkItem)
        HeapFree(GetProcessHeap(), 0x00, pWorkItem);
    return bResult;
}

