#include <windows.h>
#include <stdlib.h>
#include <iostream>
#include "session_notification.h"
using namespace std;

static BOOL WINAPI (*WTSRegisterSessionNotification)(HWND, DWORD) = NULL;
static BOOL WINAPI (*WTSUnRegisterSessionNotification)(HWND) = NULL;
static HMODULE wtsapi = NULL;

void session_notification_enable(HWND hWnd)
{
    if (wtsapi == NULL)
    {
        wtsapi = LoadLibrary("wtsapi32.dll");
        if (wtsapi == NULL)
        {
            cerr << "Could not load wtsapi32.dll!" << endl;
            exit(2);
        }
    }
    if (WTSRegisterSessionNotification == NULL)
    {
        WTSRegisterSessionNotification = (BOOL WINAPI (*)(HWND, DWORD))
                GetProcAddress(wtsapi, "WTSRegisterSessionNotification");
        if (WTSRegisterSessionNotification == NULL)
        {
            cerr << "Could not obtain WTSRegisterSessionNotification address!"
                 << endl;
            exit(3);
        }
    }

    if (WTSRegisterSessionNotification(hWnd, NOTIFY_FOR_ALL_SESSIONS)
            != TRUE)
    {
        cerr << "WTSRegisterSessionNotification() error: "
             << GetLastError() << endl;
    }
}

void session_notification_disable(HWND hWnd)
{
    if (wtsapi == NULL)
    {
        wtsapi = LoadLibrary("wtsapi32.dll");
        if (wtsapi == NULL)
        {
            cerr << "Could not load wtsapi32.dll!" << endl;
            exit(2);
        }
    }
    if (WTSUnRegisterSessionNotification == NULL)
    {
        WTSUnRegisterSessionNotification = (BOOL WINAPI (*)(HWND))
                GetProcAddress(wtsapi, "WTSUnRegisterSessionNotification");
        if (WTSUnRegisterSessionNotification == NULL)
        {
            cerr << "Could not obtain WTSUnRegisterSessionNotification address!"
                 << endl;
            exit(3);
        }
    }
    if (WTSUnRegisterSessionNotification(hWnd) != TRUE)
    {
        cerr << "WTSUnRegisterSessionNotification() error: "
             << GetLastError() << endl;
    }
}