Display Power Off Utility Code

#include "stdafx.h"

#include "Resource.h"

#include "App.h"
#include "Config.h"
#include "Engine.h"
#include "Options.h"
#include "Platform.h"
#include "SettingsDialog.h"
#include "NotificationAreaIcon.h"
#include "UserTaskDialog.h"
#include "WindowsStartup.h"
#include "Workstation.h"

constexpr auto notificationAreaIconId	  = 1;
constexpr auto notificationAreaIconIconId = IDI_APP_ICON;

constexpr auto WM_NOTIFICATION_AREA_ICON  = WM_APP + 100;

namespace
{
	UINT WM_TASKBAR_CREATED = UINT_MAX;

	CharBuffer<ARRAYSIZE(NOTIFYICONDATA::szInfo)-1> szInfo;

	[[nodiscard]] inline NOTIFYICONDATA GetNotifyIconData(HWND hwnd)
	{
#pragma warning(suppress: 26486)
		DECLARE_CBSTRUCT(NOTIFYICONDATA, notifyIconData);

		notifyIconData.uVersion = NOTIFYICON_VERSION_4;

		notifyIconData.hWnd	= hwnd;
		notifyIconData.uID	= notificationAreaIconId;

		return notifyIconData;
	}

	void LoadTip(NOTIFYICONDATA& notifyIconData)
	{
		CharBuffer<ARRAYSIZE(notifyIconData.szTip)-1> szFormat;
		CharBuffer<szFormat> szText;

		App::LoadString(IDS_NOTIFICATION_AREA_ICON_TIP_FORMAT, szFormat);
		App::LoadString(Engine::LaunchScreensaverByDefault() ? IDS_NOTIFICATION_AREA_ICON_TIP_SCREENSAVER : IDS_NOTIFICATION_AREA_ICON_TIP_TURNOFF, szText);

		SecureZeroMemory(notifyIconData.szTip, sizeof(notifyIconData.szTip));

		VERIFY_SPRINTF(wnsprintf(notifyIconData.szTip, szFormat, szFormat, szText.Chars));
	}
}

bool NotificationAreaIcon::Add(HWND hwnd)
{
	VERIFY(Workstation::RegisterSessionNotification(hwnd));

	auto notifyIconData = GetNotifyIconData(hwnd);

	notifyIconData.uFlags			= NIF_MESSAGE | NIF_ICON | NIF_TIP;
	notifyIconData.uCallbackMessage	= WM_NOTIFICATION_AREA_ICON;
	notifyIconData.hIcon			= Platform::LoadIconBest(App::GetInstance(), notificationAreaIconIconId, true);

#pragma warning(suppress: 26486)
	LoadTip(notifyIconData);

	const auto result = Shell_NotifyIcon(NIM_ADD, &notifyIconData);
	ASSERT(result && "Shell_NotifyIcon NIM_ADD");

	VERIFY(DestroyIcon(notifyIconData.hIcon));

	return result;
}

void NotificationAreaIcon::ApplySettings(HWND hwnd, bool forceConfig)
{
	const auto setNotificationAreaIconTip = [=]()
	{
		auto notifyIconData = GetNotifyIconData(hwnd);

		notifyIconData.uFlags = NIF_TIP;

#pragma warning(suppress: 26486)
		LoadTip(notifyIconData);

		const auto result = Shell_NotifyIcon(NIM_MODIFY, &notifyIconData);
		ASSERT(result && "Shell_NotifyIcon NIM_MODIFY");
	};

	Config::Read(forceConfig);

	setNotificationAreaIconTip();
}

void NotificationAreaIcon::OperationStarted(HWND hwnd, bool displayOff)
{
	if (!Options::NotificationAreaIcon())
	{
		return;
	}

	auto notifyIconData = GetNotifyIconData(hwnd);

	notifyIconData.uFlags		= NIF_INFO | NIF_SHOWTIP | NIF_REALTIME | NIF_TIP;
	notifyIconData.dwInfoFlags	= NIIF_USER | NIIF_LARGE_ICON;
	notifyIconData.hBalloonIcon	= Platform::LoadIconBest(App::GetInstance(), notificationAreaIconIconId);

	CharBuffer<ARRAYSIZE(notifyIconData.szInfoTitle)-1> szInfoTitle;

	App::LoadString(IDS_APP_TITLE, szInfoTitle);
	App::LoadString(displayOff ? IDS_DISPLAY_OFF_PROGRESS : IDS_SCREENSAVER_PROGRESS, szInfo);

	VERIFY_STR(StringCchCopy(notifyIconData.szInfoTitle, szInfoTitle, szInfoTitle));
	VERIFY_STR(StringCchCopy(notifyIconData.szInfo, szInfo, szInfo));
	VERIFY_STR(StringCchCopy(notifyIconData.szTip, ARRAYSIZE(notifyIconData.szTip), szInfo));

#pragma warning(suppress: 26486)
	const auto result = Shell_NotifyIcon(NIM_MODIFY, &notifyIconData);

	VERIFY(DestroyIcon(notifyIconData.hBalloonIcon));

	ASSERT(result && "Shell_NotifyIcon NIM_MODIFY");
}

void NotificationAreaIcon::OperationEnded(HWND hwnd)
{
	if (!Options::NotificationAreaIcon())
	{
		return;
	}

	ApplySettings(hwnd, false);
}

namespace
{
	void DeleteNotificationAreaIcon(HWND hwnd)
	{
		auto notifyIconData = GetNotifyIconData(hwnd);
#pragma warning(suppress: 26486)
		Shell_NotifyIcon(NIM_DELETE, &notifyIconData);
	}

	void AdjustMenu(HMENU hMenu)
	{
		const auto whenNotBusy = [=]()
		{
			Config::Read(true);

			const auto launchScreensaver = Engine::LaunchScreensaverByDefault();

			if (launchScreensaver)
			{
				VERIFY(Platform::MoveMenuItem(hMenu, 1, 0));
			}

			CharBuffer<256> cszText;
			App::LoadString(launchScreensaver ? IDS_SCREENSAVER_ON_LOCK : IDS_TURN_OFF_ON_LOCK, cszText);

			VERIFY(Platform::SetMenuItemText(hMenu, ID_TRIGGER_ON_WORKSTATION_LOCK, cszText));
		};

		const auto whenBusy = [=]()
		{
			VERIFY(Platform::SetMenuItemText(hMenu, ID_IN_PROGRESS, szInfo));
		};

		if (Engine::Busy())
		{
			whenBusy();
		}
		else
		{
			whenNotBusy();
		}

		VERIFY(SetMenuDefaultItem(hMenu, 0, TRUE));

		const auto checkMenuItem = [=](UINT id, bool check)
		{
			CheckMenuItem(hMenu, id, MF_BYCOMMAND | (check ? MF_CHECKED : MF_UNCHECKED));
		};

		checkMenuItem(ID_RUN_AT_WINDOWS_STARTUP,		WindowsStartup::Enabled());
		checkMenuItem(ID_TRIGGER_ON_WORKSTATION_LOCK,	Config::TriggerOnLockWorkStation());
	}

	void DisplayContextMenu(HWND hwnd)
	{
		constexpr auto getMenuId = []() { return Engine::Busy() ? IDR_IN_PROGRESS_POPUP_MENU : IDR_POPUP_MENU; };

		const auto hMenu = LoadMenu(App::GetInstance(), MAKEINTRESOURCE(getMenuId()));
		ASSERT(hMenu);

		const auto hMenuTrackPopup = GetSubMenu(hMenu, 0);
		ASSERT(hMenuTrackPopup);

		VERIFY(SetForegroundWindow(hwnd));

		POINT point;
		VERIFY(GetCursorPos(&point));

		AdjustMenu(hMenuTrackPopup);
		VERIFY(TrackPopupMenu(hMenuTrackPopup, GetSystemMetrics(SM_MENUDROPALIGNMENT) | TPM_RIGHTBUTTON | TPM_LEFTBUTTON, point.x, point.y, 0, hwnd, nullptr));

		VERIFY(PostMessage(hwnd, WM_NULL, 0, 0));

		VERIFY(DestroyMenu(hMenuTrackPopup));
		VERIFY(DestroyMenu(hMenu));
	}

	BOOL OnCreate(HWND, LPCREATESTRUCT)
	{
		const auto wmTaskbarCreated = RegisterWindowMessage(TEXT("TaskbarCreated"));
		ASSERT(wmTaskbarCreated);

		if (wmTaskbarCreated)
		{
			WM_TASKBAR_CREATED = wmTaskbarCreated;
		}

		return TRUE;
	}

	void OnCommand(HWND hwnd, int id, HWND, UINT)
	{
		const auto onTriggerOnWorkStationLock = []()
		{
			const auto set = !Config::TriggerOnLockWorkStation();
			SettingsDialog::SetTriggerOnWorkStationLockCheck(set);
			Config::WriteTriggerOnLockWorkStation(set);
		};

		const auto onRunAtWindowsStartup = []()
		{
			const auto set = !WindowsStartup::Enabled();
			SettingsDialog::SetRunAtStartupCheck(set);
			WindowsStartup::Enable(set);
		};

		HANDLE_COMMAND(ID_TURN_DISPLAY_OFF,				Engine::TurnOffDisplay(hwnd));
		HANDLE_COMMAND(ID_LAUNCH_SCREENSAVER,			Engine::LaunchScreensaver(hwnd));
		HANDLE_COMMAND(ID_SETTINGS,						Platform::InitCommonControls(); SettingsDialog::Show(hwnd));
		HANDLE_COMMAND(ID_RUN_AT_WINDOWS_STARTUP,		onRunAtWindowsStartup());
		HANDLE_COMMAND(ID_TRIGGER_ON_WORKSTATION_LOCK,	onTriggerOnWorkStationLock());
		HANDLE_COMMAND(ID_ABOUT,						UserTaskDialog::Help());
		HANDLE_COMMAND(ID_EXIT,							Workstation::UnregisterSessionNotification(hwnd); DeleteNotificationAreaIcon(hwnd); Platform::CloseWindow(hwnd));
	}

	LRESULT OnNotify(HWND hwnd, UINT uMsg)
	{
		HANDLE_APP_MSG(WM_LBUTTONDBLCLK,				Engine::ExecuteDefault(hwnd));
		HANDLE_APP_MSG(WM_CONTEXTMENU,					DisplayContextMenu(hwnd));
		HANDLE_APP_MSG(WM_RBUTTONUP,					DisplayContextMenu(hwnd));

		return 0;
	}
}

LRESULT NotificationAreaIcon::Process(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
	HANDLE_APP_MSG(WM_TASKBAR_CREATED,					DeleteNotificationAreaIcon(hwnd); [[maybe_unused]] const auto _ = NotificationAreaIcon::Add(hwnd));
	HANDLE_APP_MSG(WM_NOTIFICATION_AREA_ICON,			OnNotify(hwnd, LOWORD(lParam)));

#pragma warning (suppress: 26818)
	switch (uMsg)
	{
		HANDLE_MSG(hwnd, WM_CREATE,		OnCreate);
		HANDLE_MSG(hwnd, WM_COMMAND,	OnCommand);
		HANDLE_MSG(hwnd, WM_SETFOCUS,	[](HWND, HWND) { UserTaskDialog::MakeForeground(); SettingsDialog::MakeForeground(); });
	}

	return 0;
}