#include <Ldk/Windows.h>

DRIVER_INITIALIZE DriverEntry;
DRIVER_UNLOAD DriverUnload;

typedef struct _COND_TEST_THREAD_PARAM {
    PCSTR name;
    BOOL ready;
    LONG processed;
    CONDITION_VARIABLE cv;
    CRITICAL_SECTION cs;
} COND_TEST_THREAD_PARAM, *PCOND_TEST_THREAD_PARAM;

DWORD
WINAPI
CondTestThreadProc (
    LPVOID lpThreadParameter
    )
{
    PCOND_TEST_THREAD_PARAM param = (PCOND_TEST_THREAD_PARAM)lpThreadParameter;

    PAGED_CODE();

    DbgPrint("%s - %d - start\n", param->name, GetCurrentThreadId());

    EnterCriticalSection(&param->cs);
    while (!param->ready) {
        SleepConditionVariableCS(&param->cv, &param->cs, INFINITE);
    }

    DbgPrint("%s - %d - processed (%d)\n", param->name, GetCurrentThreadId(), InterlockedIncrement(&param->processed));

    LeaveCriticalSection(&param->cs);

    DbgPrint("%s - %d - end\n", param->name, GetCurrentThreadId());

    WakeConditionVariable(&param->cv);
    return 0;
}

NTSTATUS
DriverEntry (
    _In_ PDRIVER_OBJECT DriverObject,
    _In_ PUNICODE_STRING RegistryPath
    )
{
    PAGED_CODE();
    UNREFERENCED_PARAMETER(RegistryPath);

    COND_TEST_THREAD_PARAM param;
    param.name = "Test";
    param.ready = FALSE;
    param.processed = 0;
    InitializeCriticalSection(&param.cs);
    InitializeConditionVariable(&param.cv);
    HANDLE threads[5];
    for (int i = 0; i < 5; ++i) {
        threads[i] = CreateThread(NULL, 0, CondTestThreadProc, &param, 0, NULL);
    }

    EnterCriticalSection(&param.cs);
    param.ready = TRUE;
    LeaveCriticalSection(&param.cs);
    WakeAllConditionVariable(&param.cv);

    EnterCriticalSection(&param.cs);
    while (param.processed < 5) {
        SleepConditionVariableCS(&param.cv, &param.cs, INFINITE);
    }
    WakeAllConditionVariable(&param.cv);
    LeaveCriticalSection(&param.cs);

    WaitForMultipleObjects(5, threads, TRUE, INFINITE);
    DeleteCriticalSection(&param.cs);

    for (int i = 0; i < 5; ++i) {
        if (threads[i]) {
            CloseHandle(threads[i]);
        }
    }

    DriverObject->DriverUnload = DriverUnload;
    return STATUS_SUCCESS;
}

VOID
DriverUnload (
    _In_ PDRIVER_OBJECT driverObject
    )
{
    PAGED_CODE();
    UNREFERENCED_PARAMETER(driverObject);
}