tray_icon.cpp

#include "pch.h"
#include "Generated files/resource.h"
#include "settings_window.h"
#include "tray_icon.h"
#include "general_settings.h"
#include "centralized_hotkeys.h"
#include "centralized_kb_hook.h"
#include "quick_access_host.h"
#include "hotkey_conflict_detector.h"
#include "trace.h"
#include <Windows.h>

#include <common/utils/resources.h>
#include <common/utils/window.h>
#include <common/version/version.h>
#include <common/logger/logger.h>
#include <common/utils/elevation.h>
#include <common/Themes/theme_listener.h>
#include <common/Themes/theme_helpers.h>
#include "bug_report.h"
#include <common/updating/updateState.h>

namespace
{
    HWND tray_icon_hwnd = NULL;

    enum
    {
        wm_icon_notify = WM_APP,
        wm_run_on_main_ui_thread,
    };

    // Contains the Windows Message for taskbar creation.
    UINT wm_taskbar_restart = 0;

    NOTIFYICONDATAW tray_icon_data;
    bool tray_icon_created = false;

    // Set when Windows confirms that the full session is ending so WM_DESTROY
    // can skip cross-process cleanup the OS is already performing in parallel.
    bool g_system_session_ending = false;

    bool about_box_shown = false;

    HMENU h_menu = nullptr;
    HMENU h_sub_menu = nullptr;
    bool double_click_timer_running = false;
    bool double_clicked = false;
    POINT tray_icon_click_point;
    std::optional<bool> last_quick_access_state; // Track the last known Quick Access state

    static ThemeListener theme_listener;
    static bool theme_adaptive_enabled = false;
    static bool update_available = false;

    void log_session_end_message(UINT message, WPARAM wparam, LPARAM lparam)
    {
        if (message != WM_QUERYENDSESSION && message != WM_ENDSESSION)
        {
            return;
        }

        const auto flags = static_cast<DWORD>(lparam);
        const bool close_app = (flags & ENDSESSION_CLOSEAPP) != 0;
        const bool critical = (flags & ENDSESSION_CRITICAL) != 0;
        const bool logoff = (flags & ENDSESSION_LOGOFF) != 0;
        const wchar_t* reason = L"shutdown_or_restart";
        if (close_app)
        {
            reason = L"restart_manager_close_app";
        }
        else if (logoff)
        {
            reason = L"logoff";
        }

        if (message == WM_QUERYENDSESSION)
        {
            Logger::info(L"Runner received WM_QUERYENDSESSION: reason={}, flags={:#010x}, close_app={}, critical={}, logoff={}",
                         reason,
                         flags,
                         close_app,
                         critical,
                         logoff);
        }
        else
        {
            Logger::info(L"Runner received WM_ENDSESSION: session_ending={}, reason={}, flags={:#010x}, close_app={}, critical={}, logoff={}",
                         wparam != FALSE,
                         reason,
                         flags,
                         close_app,
                         critical,
                         logoff);
        }
    }
}

// Struct to fill with callback and the data. The window_proc is responsible for cleaning it.
struct run_on_main_ui_thread_msg
{
    main_loop_callback_function _callback;
    PVOID data;
};

bool dispatch_run_on_main_ui_thread(main_loop_callback_function _callback, PVOID data)
{
    if (tray_icon_hwnd == NULL)
    {
        return false;
    }
    struct run_on_main_ui_thread_msg* wnd_msg = new struct run_on_main_ui_thread_msg();
    wnd_msg->_callback = _callback;
    wnd_msg->data = data;

    PostMessage(tray_icon_hwnd, wm_run_on_main_ui_thread, 0, reinterpret_cast<LPARAM>(wnd_msg));

    return true;
}

void change_menu_item_text(const UINT item_id, wchar_t* new_text)
{
    MENUITEMINFOW menuitem = { .cbSize = sizeof(MENUITEMINFOW), .fMask = MIIM_TYPE | MIIM_DATA };
    GetMenuItemInfoW(h_menu, item_id, false, &menuitem);
    menuitem.dwTypeData = new_text;
    SetMenuItemInfoW(h_menu, item_id, false, &menuitem);
}

void open_quick_access_flyout_window()
{
    QuickAccessHost::show();
}

void handle_tray_command(HWND window, const WPARAM command_id, LPARAM lparam)
{
    switch (command_id)
    {
    case ID_SETTINGS_MENU_COMMAND:
    {
        std::wstring settings_window{ winrt::to_hstring(ESettingsWindowNames_to_string(static_cast<ESettingsWindowNames>(lparam))) };
        open_settings_window(settings_window);
    }
    break;
    case ID_CLOSE_MENU_COMMAND:
        if (h_menu)
        {
            DestroyMenu(h_menu);
        }
        DestroyWindow(window);
        break;
    case ID_ABOUT_MENU_COMMAND:
        if (!about_box_shown)
        {
            about_box_shown = true;
            std::wstring about_msg = L"PowerToys\nVersion " + get_product_version() + L"\n\xa9 2019 Microsoft Corporation";
            MessageBoxW(nullptr, about_msg.c_str(), L"About PowerToys", MB_OK);
            about_box_shown = false;
        }
        break;
    case ID_REPORT_BUG_COMMAND:
    {
        launch_bug_report();
        break;
    }

    case ID_DOCUMENTATION_MENU_COMMAND:
    {
        RunNonElevatedEx(L"https://aka.ms/PowerToysOverview", L"", L"");
        break;
    }
    case ID_QUICK_ACCESS_MENU_COMMAND:
    {
        open_quick_access_flyout_window();
        break;
    }
    case ID_UPDATE_MENU_COMMAND:
    {
        open_settings_window(std::wstring{ L"Overview" });
        break;
    }
    }
}

void click_timer_elapsed()
{
    double_click_timer_running = false;
    if (!double_clicked)
    {
        // Log telemetry for single click (confirmed it's not a double click)
        Trace::TrayIconLeftClick(get_general_settings().enableQuickAccess);

        if (get_general_settings().enableQuickAccess)
        {
            open_quick_access_flyout_window();
        }
        else
        {
            open_settings_window(std::nullopt);
        }
    }
}

LRESULT __stdcall tray_icon_window_proc(HWND window, UINT message, WPARAM wparam, LPARAM lparam)
{
    log_session_end_message(message, wparam, lparam);

    LRESULT session_end_result = 0;
    if (handle_stateless_session_end_message(window, message, wparam, lparam, session_end_result, &g_system_session_ending))
    {
        return session_end_result;
    }

    switch (message)
    {
    case WM_HOTKEY:
    {
        // We use the tray icon WndProc to avoid creating a dedicated window just for this message.
        const auto modifiersMask = LOWORD(lparam);
        const auto vkCode = HIWORD(lparam);
        Logger::trace(L"On {} hotkey", CentralizedHotkeys::ToWstring({ modifiersMask, vkCode }));
        CentralizedHotkeys::PopulateHotkey({ modifiersMask, vkCode });
        break;
    }
    case WM_CREATE:
        if (wm_taskbar_restart == 0)
        {
            tray_icon_hwnd = window;
            wm_taskbar_restart = RegisterWindowMessageW(L"TaskbarCreated");
        }
        break;
    case WM_DESTROY:
        // On OS-initiated shutdown skip cross-process cleanup: the shell is tearing
        // down and close_settings_window() blocks up to 1.5s waiting on
        // PowerToys.Settings.exe, which the OS is reaping in parallel. That wait, plus
        // Shell_NotifyIcon during explorer teardown, would burn the limited quiesce
        // budget and trip APPLICATION_HANG_QUIESCE. PostQuitMessage alone unwinds the
        // loop in milliseconds. User-initiated Exit and Restart Manager
        // ENDSESSION_CLOSEAPP requests keep the full graceful cleanup.
        Logger::info(L"Runner WM_DESTROY, system_session_ending={}", g_system_session_ending);
        if (!g_system_session_ending)
        {
            if (tray_icon_created)
            {
                Shell_NotifyIcon(NIM_DELETE, &tray_icon_data);
                tray_icon_created = false;
            }
            close_settings_window();
        }
        PostQuitMessage(0);
        break;
    case WM_CLOSE:
        DestroyWindow(window);
        break;
    case WM_COMMAND:
        handle_tray_command(window, wparam, lparam);
        break;
    // Shell_NotifyIcon can fail when we invoke it during the time explorer.exe isn't present/ready to handle it.
    // We'll also never receive wm_taskbar_restart message if the first call to Shell_NotifyIcon failed, so we use
    // WM_WINDOWPOSCHANGING which is always received on explorer startup sequence.
    case WM_WINDOWPOSCHANGING:
    {
        if (!tray_icon_created)
        {
            tray_icon_created = Shell_NotifyIcon(NIM_ADD, &tray_icon_data) == TRUE;
        }
        break;
    }
    default:
        if (message == wm_icon_notify)
        {
            switch (lparam)
            {
            case WM_RBUTTONUP:
            case WM_CONTEXTMENU:
            {
                bool quick_access_enabled = get_general_settings().enableQuickAccess;

                // Log telemetry
                Trace::TrayIconRightClick(quick_access_enabled);

                // Reload menu if Quick Access state has changed or is first time
                if (h_menu && (!last_quick_access_state.has_value() || quick_access_enabled != last_quick_access_state.value()))
                {
                    DestroyMenu(h_menu);
                    h_menu = nullptr;
                    h_sub_menu = nullptr;
                }

                last_quick_access_state = quick_access_enabled;

                if (!h_menu)
                {
                    h_menu = LoadMenu(reinterpret_cast<HINSTANCE>(&__ImageBase), MAKEINTRESOURCE(ID_TRAY_MENU));
                }
                if (h_menu)
                {
                    static std::wstring settings_menuitem_label = GET_RESOURCE_STRING(IDS_SETTINGS_MENU_TEXT);
                    static std::wstring settings_menuitem_label_leftclick = GET_RESOURCE_STRING(IDS_SETTINGS_MENU_TEXT_LEFTCLICK);
                    static std::wstring close_menuitem_label = GET_RESOURCE_STRING(IDS_CLOSE_MENU_TEXT);
                    static std::wstring submit_bug_menuitem_label = GET_RESOURCE_STRING(IDS_SUBMIT_BUG_TEXT);
                    static std::wstring documentation_menuitem_label = GET_RESOURCE_STRING(IDS_DOCUMENTATION_MENU_TEXT);
                    static std::wstring quick_access_menuitem_label = GET_RESOURCE_STRING(IDS_QUICK_ACCESS_MENU_TEXT);

                    // Update Settings menu text based on Quick Access state
                    if (quick_access_enabled)
                    {
                        change_menu_item_text(ID_SETTINGS_MENU_COMMAND, settings_menuitem_label.data());
                    }
                    else
                    {
                        change_menu_item_text(ID_SETTINGS_MENU_COMMAND, settings_menuitem_label_leftclick.data());
                    }

                    change_menu_item_text(ID_CLOSE_MENU_COMMAND, close_menuitem_label.data());
                    change_menu_item_text(ID_REPORT_BUG_COMMAND, submit_bug_menuitem_label.data());
                    bool bug_report_disabled = is_bug_report_running();
                    EnableMenuItem(h_sub_menu, ID_REPORT_BUG_COMMAND, MF_BYCOMMAND | (bug_report_disabled ? MF_GRAYED : MF_ENABLED));
                    change_menu_item_text(ID_DOCUMENTATION_MENU_COMMAND, documentation_menuitem_label.data());
                    change_menu_item_text(ID_QUICK_ACCESS_MENU_COMMAND, quick_access_menuitem_label.data());

                    // Hide or show Quick Access menu item based on setting
                    if (!h_sub_menu)
                    {
                        h_sub_menu = GetSubMenu(h_menu, 0);
                    }
                    if (!quick_access_enabled)
                    {
                        // Remove Quick Access menu item when disabled
                        DeleteMenu(h_sub_menu, ID_QUICK_ACCESS_MENU_COMMAND, MF_BYCOMMAND);
                    }
                }
                if (!h_sub_menu)
                {
                    h_sub_menu = GetSubMenu(h_menu, 0);
                }

                // Dynamically add/remove "Update available" menu item and its separator
                DeleteMenu(h_sub_menu, ID_UPDATE_MENU_COMMAND, MF_BYCOMMAND);
                // Remove the separator right after the update item (position 0 after deletion)
                if (GetMenuItemCount(h_sub_menu) > 0)
                {
                    MENUITEMINFOW mii = { .cbSize = sizeof(mii), .fMask = MIIM_FTYPE };
                    if (GetMenuItemInfoW(h_sub_menu, 0, TRUE, &mii) && (mii.fType & MFT_SEPARATOR))
                    {
                        DeleteMenu(h_sub_menu, 0, MF_BYPOSITION);
                    }
                }
                if (update_available)
                {
                    InsertMenuW(h_sub_menu, 0, MF_BYPOSITION | MF_STRING, ID_UPDATE_MENU_COMMAND, GET_RESOURCE_STRING(IDS_UPDATE_AVAILABLE_MENU_TEXT).c_str());
                    InsertMenuW(h_sub_menu, 1, MF_BYPOSITION | MF_SEPARATOR, 0, nullptr);
                }

                POINT mouse_pointer;
                GetCursorPos(&mouse_pointer);
                SetForegroundWindow(window); // Needed for the context menu to disappear.
                TrackPopupMenu(h_sub_menu, TPM_CENTERALIGN | TPM_BOTTOMALIGN, mouse_pointer.x, mouse_pointer.y, 0, window, nullptr);
                break;
            }
            case WM_LBUTTONUP:
            {
                // ignore event if this is the second click of a double click
                if (!double_click_timer_running)
                {
                    // start timer for detecting single or double click
                    double_click_timer_running = true;
                    double_clicked = false;

                    UINT doubleClickTime = GetDoubleClickTime();
                    std::thread([doubleClickTime]() {
                        std::this_thread::sleep_for(std::chrono::milliseconds(doubleClickTime));
                        click_timer_elapsed();
                    }).detach();
                }
                break;
            }
            case WM_LBUTTONDBLCLK:
            {
                // Log telemetry
                Trace::TrayIconDoubleClick(get_general_settings().enableQuickAccess);

                double_clicked = true;
                open_settings_window(std::nullopt);
                break;
            }
            break;
            }
        }
        else if (message == wm_run_on_main_ui_thread)
        {
            if (lparam != NULL)
            {
                struct run_on_main_ui_thread_msg* msg = reinterpret_cast<struct run_on_main_ui_thread_msg*>(lparam);
                msg->_callback(msg->data);
                delete msg;
                lparam = NULL;
            }
            break;
        }
        else if (message == wm_taskbar_restart)
        {
            tray_icon_created = Shell_NotifyIcon(NIM_ADD, &tray_icon_data) == TRUE;
            break;
        }
    }
    return DefWindowProc(window, message, wparam, lparam);
}

static HICON get_icon(Theme theme)
{
    std::wstring icon_path = get_module_folderpath();
    if (theme == Theme::Dark)
    {
        icon_path += update_available ? L"\\svgs\\PowerToysWhiteUpdate.ico" : L"\\svgs\\PowerToysWhite.ico";
    }
    else
    {
        icon_path += update_available ? L"\\svgs\\PowerToysDarkUpdate.ico" : L"\\svgs\\PowerToysDark.ico";
    }
    Logger::trace(L"get_icon: Loading icon from path: {}", icon_path);

    HICON icon = static_cast<HICON>(LoadImage(NULL,
                                        icon_path.c_str(),
                                        IMAGE_ICON,
                                        0,
                                        0,
                                        LR_LOADFROMFILE | LR_DEFAULTSIZE | LR_SHARED));
    if (!icon)
    {
        Logger::warn(L"get_icon: Failed to load icon from {}, error: {}", icon_path, GetLastError());
    }
    return icon;
}


static void handle_theme_change()
{
    if (theme_adaptive_enabled)
    {
        tray_icon_data.hIcon = get_icon(ThemeHelpers::GetSystemTheme());
        Shell_NotifyIcon(NIM_MODIFY, &tray_icon_data);
    }
}

void update_bug_report_menu_status(bool isRunning)
{
    if (h_sub_menu != nullptr)
    {
        EnableMenuItem(h_sub_menu, ID_REPORT_BUG_COMMAND, MF_BYCOMMAND | (isRunning ? MF_GRAYED : MF_ENABLED));
    }
}

void start_tray_icon(bool isProcessElevated, bool theme_adaptive)
{
    theme_adaptive_enabled = theme_adaptive;
    auto h_instance = reinterpret_cast<HINSTANCE>(&__ImageBase);

    // Check if an update is available at startup
    auto state = UpdateState::read();
    update_available = (state.state == UpdateState::readyToDownload || state.state == UpdateState::readyToInstall);

    HICON const icon = theme_adaptive
                           ? get_icon(ThemeHelpers::GetSystemTheme())
                           : LoadIcon(h_instance, MAKEINTRESOURCE(update_available ? APPICON_UPDATE : APPICON));
    if (icon)
    {
        UINT id_tray_icon = 1;

        WNDCLASS wc = {};
        wc.hCursor = LoadCursor(nullptr, IDC_ARROW);
        wc.hInstance = h_instance;
        wc.lpszClassName = pt_tray_icon_window_class;
        wc.style = CS_HREDRAW | CS_VREDRAW;
        wc.lpfnWndProc = tray_icon_window_proc;
        wc.hIcon = icon;
        RegisterClass(&wc);
        auto hwnd = CreateWindowW(wc.lpszClassName,
                                  pt_tray_icon_window_class,
                                  WS_OVERLAPPEDWINDOW | WS_POPUP,
                                  CW_USEDEFAULT,
                                  CW_USEDEFAULT,
                                  CW_USEDEFAULT,
                                  CW_USEDEFAULT,
                                  nullptr,
                                  nullptr,
                                  wc.hInstance,
                                  nullptr);
        WINRT_VERIFY(hwnd);
        CentralizedHotkeys::RegisterWindow(hwnd);
        CentralizedKeyboardHook::RegisterWindow(hwnd);
        memset(&tray_icon_data, 0, sizeof(tray_icon_data));
        tray_icon_data.cbSize = sizeof(tray_icon_data);
        tray_icon_data.hIcon = icon;
        tray_icon_data.hWnd = hwnd;
        tray_icon_data.uID = id_tray_icon;
        tray_icon_data.uCallbackMessage = wm_icon_notify;

        std::wstringstream pt_version_tooltip_stream;
        if (isProcessElevated)
        {
            pt_version_tooltip_stream << GET_RESOURCE_STRING(IDS_TRAY_ICON_ADMIN_TOOLTIP) << L": ";
        }

        pt_version_tooltip_stream << L"PowerToys " << get_product_version() << '\0';
        std::wstring pt_version_tooltip = pt_version_tooltip_stream.str();
        wcscpy_s(tray_icon_data.szTip, sizeof(tray_icon_data.szTip) / sizeof(WCHAR), pt_version_tooltip.c_str());
        tray_icon_data.uFlags = NIF_ICON | NIF_TIP | NIF_MESSAGE;
        ChangeWindowMessageFilterEx(hwnd, WM_COMMAND, MSGFLT_ALLOW, nullptr);

        tray_icon_created = Shell_NotifyIcon(NIM_ADD, &tray_icon_data) == TRUE;
        theme_listener.AddSystemThemeChangedHandler(&handle_theme_change);

        // Register callback to update bug report menu item status
        BugReportManager::instance().register_callback([](bool isRunning) {
            dispatch_run_on_main_ui_thread([](PVOID data) {
                bool* running = static_cast<bool*>(data);
                update_bug_report_menu_status(*running);
                delete running;
            },
                                           new bool(isRunning));
        });
    }
}

void set_tray_icon_visible(bool shouldIconBeVisible)
{
    tray_icon_data.uFlags |= NIF_STATE;
    tray_icon_data.dwStateMask = NIS_HIDDEN;
    tray_icon_data.dwState = shouldIconBeVisible ? 0 : NIS_HIDDEN;
    Shell_NotifyIcon(NIM_MODIFY, &tray_icon_data);
}

void set_tray_icon_update_available(bool available)
{
    if (update_available == available)
    {
        return;
    }

    update_available = available;
    Logger::info(L"set_tray_icon_update_available: update_available={}", update_available);

    if (theme_adaptive_enabled)
    {
        tray_icon_data.hIcon = get_icon(ThemeHelpers::GetSystemTheme());
    }
    else
    {
        auto h_instance = reinterpret_cast<HINSTANCE>(&__ImageBase);
        tray_icon_data.hIcon = LoadIcon(h_instance, MAKEINTRESOURCE(available ? APPICON_UPDATE : APPICON));
    }

    Shell_NotifyIcon(NIM_MODIFY, &tray_icon_data);
}

void set_tray_icon_theme_adaptive(bool theme_adaptive)
{
    Logger::info(L"set_tray_icon_theme_adaptive: Called with theme_adaptive={}, current theme_adaptive_enabled={}",
                 theme_adaptive, theme_adaptive_enabled);

    auto h_instance = reinterpret_cast<HINSTANCE>(&__ImageBase);
    HICON icon = nullptr;

    if (theme_adaptive)
    {
        icon = get_icon(ThemeHelpers::GetSystemTheme());
        if (!icon)
        {
            Logger::warn(L"set_tray_icon_theme_adaptive: Failed to load theme adaptive icon, falling back to default");
        }
    }

    // If not requesting adaptive icon, or if adaptive icon failed to load, use default icon
    if (!icon)
    {
        icon = LoadIcon(h_instance, MAKEINTRESOURCE(update_available ? APPICON_UPDATE : APPICON));
        if (theme_adaptive && icon)
        {
            // We requested adaptive but had to fall back, so update the flag
            theme_adaptive = false;
            Logger::info(L"set_tray_icon_theme_adaptive: Using default icon as fallback");
        }
    }

    theme_adaptive_enabled = theme_adaptive;

    if (icon)
    {
        tray_icon_data.hIcon = icon;
        BOOL result = Shell_NotifyIcon(NIM_MODIFY, &tray_icon_data);
        Logger::info(L"set_tray_icon_theme_adaptive: Icon updated, theme_adaptive_enabled={}, Shell_NotifyIcon result={}",
                     theme_adaptive_enabled, result);
    }
    else
    {
        Logger::error(L"set_tray_icon_theme_adaptive: Failed to load any icon");
    }
}

void stop_tray_icon()
{
    if (tray_icon_created)
    {
        // Clear bug report callbacks
        BugReportManager::instance().clear_callbacks();
        SendMessage(tray_icon_hwnd, WM_CLOSE, 0, 0);
    }
}

bool is_system_session_ending()
{
    return g_system_session_ending;
}
void update_quick_access_hotkey(bool enabled, PowerToysSettings::HotkeyObject hotkey)
{
    static PowerToysSettings::HotkeyObject current_hotkey;
    static bool is_registered = false;
    auto& hkmng = HotkeyConflictDetector::HotkeyConflictManager::GetInstance();

    if (is_registered)
    {
        CentralizedKeyboardHook::ClearModuleHotkeys(L"QuickAccess");
        hkmng.RemoveHotkeyByModule(L"GeneralSettings");
        is_registered = false;
    }

    if (enabled && hotkey.get_code() != 0)
    {
        HotkeyConflictDetector::Hotkey hk = {
            hotkey.win_pressed(),
            hotkey.ctrl_pressed(),
            hotkey.shift_pressed(),
            hotkey.alt_pressed(),
            static_cast<unsigned char>(hotkey.get_code())
        };

        hkmng.AddHotkey(hk, L"GeneralSettings", 0, true);
        CentralizedKeyboardHook::SetHotkeyAction(L"QuickAccess", hk, []() {
            open_quick_access_flyout_window();
            return true;
        });

        current_hotkey = hotkey;
        is_registered = true;
    }
}