sherpa_server/tray_icon_win32.cpp

#ifndef _CRT_SECURE_NO_WARNINGS
#define _CRT_SECURE_NO_WARNINGS
#endif

#include "tray_icon_win32.hpp"

#ifdef _WIN32

#include <windows.h>
#include <shellapi.h>

#include <atomic>
#include <cstring>
#include <functional>
#include <utility>

#pragma comment(lib, "user32.lib")
#pragma comment(lib, "shell32.lib")

namespace sherpa {

namespace {

// WM_APP-based callback ID for our Shell_NotifyIcon. Anything in [WM_APP,
// 0xBFFF] is safe to use for app-defined messages.
constexpr UINT WM_TRAYICON = WM_APP + 1;

// Unique ID for our notification icon within the HWND. Arbitrary value.
constexpr UINT TRAY_ICON_UID = 0xC1A2;

// Right-click menu command ids.
constexpr UINT IDM_TOOLTIP = 1;  // disabled label echoing the tooltip
constexpr UINT IDM_EXIT    = 2;  // shut the helper down

HANDLE  g_thread       = nullptr;
DWORD   g_threadId     = 0;
char    g_tooltip[128] = {};

// Set inside the tray thread once NIM_ADD succeeds, so StopTrayIcon can
// tell the difference between "still starting" and "icon is live".
std::atomic<bool> g_iconLive{false};

// Optional graceful-shutdown hook, invoked on the tray thread when the user
// picks "Exit". Empty -> Exit terminates the process via ExitProcess.
std::function<void()> g_onExit;

// Handle the "Exit" menu pick (tray thread). Remove the icon immediately so
// no ghost lingers, then either hand back to the caller's shutdown hook or
// terminate outright.
void RequestExit(HWND hwnd) {
    NOTIFYICONDATAA nidDel = {};
    nidDel.cbSize = sizeof(nidDel);
    nidDel.hWnd   = hwnd;
    nidDel.uID    = TRAY_ICON_UID;
    Shell_NotifyIconA(NIM_DELETE, &nidDel);
    g_iconLive.store(false, std::memory_order_release);

    if (g_onExit) {
        // Graceful: e.g. stop the listen loop so main() returns and calls
        // StopTrayIcon(), which unwinds this message thread.
        g_onExit();
    } else {
        // No graceful hook (blocking stdin loop): terminate. The OS reaps
        // the already-removed icon and this thread.
        ExitProcess(0);
    }
}

LRESULT CALLBACK TrayWndProc(HWND hwnd, UINT msg,
                              WPARAM wParam, LPARAM lParam) {
    if (msg == WM_TRAYICON) {
        UINT event = LOWORD(lParam);
        if (event == WM_RBUTTONUP || event == WM_CONTEXTMENU) {
            // Right-click menu: a disabled label mirroring the tooltip, a
            // separator, then an "Exit" item that shuts the helper down.
            POINT pt;
            GetCursorPos(&pt);
            HMENU menu = CreatePopupMenu();
            if (menu) {
                AppendMenuA(menu, MF_STRING | MF_GRAYED, IDM_TOOLTIP, g_tooltip);
                AppendMenuA(menu, MF_SEPARATOR, 0, nullptr);
                AppendMenuA(menu, MF_STRING, IDM_EXIT, "Exit");
                // SetForegroundWindow is required so the menu dismisses
                // correctly when the user clicks outside it.
                SetForegroundWindow(hwnd);
                // TPM_RETURNCMD returns the picked id inline instead of
                // posting WM_COMMAND, so we can act on it right here.
                UINT cmd = TrackPopupMenu(menu, TPM_RIGHTBUTTON | TPM_RETURNCMD,
                                          pt.x, pt.y, 0, hwnd, NULL);
                DestroyMenu(menu);
                if (cmd == IDM_EXIT) RequestExit(hwnd);
            }
        }
        return 0;
    }
    if (msg == WM_DESTROY) {
        PostQuitMessage(0);
        return 0;
    }
    return DefWindowProcA(hwnd, msg, wParam, lParam);
}

DWORD WINAPI TrayThreadProc(LPVOID /*param*/) {
    HINSTANCE hInst = GetModuleHandleA(NULL);

    WNDCLASSEXA wc = {};
    wc.cbSize        = sizeof(wc);
    wc.lpfnWndProc   = TrayWndProc;
    wc.hInstance     = hInst;
    wc.lpszClassName = "SherpaTrayClass";
    if (!RegisterClassExA(&wc) && GetLastError() != ERROR_CLASS_ALREADY_EXISTS) {
        return 1;
    }

    HWND hwnd = CreateWindowExA(0, "SherpaTrayClass", "Sherpa TTS",
                                 0, 0, 0, 0, 0,
                                 HWND_MESSAGE,  // message-only window
                                 NULL, hInst, NULL);
    if (!hwnd) return 1;

    NOTIFYICONDATAA nid = {};
    nid.cbSize           = sizeof(nid);
    nid.hWnd             = hwnd;
    nid.uID              = TRAY_ICON_UID;
    nid.uFlags           = NIF_MESSAGE | NIF_ICON | NIF_TIP;
    nid.uCallbackMessage = WM_TRAYICON;
    nid.hIcon            = LoadIconA(NULL, IDI_APPLICATION);
    // szTip is 128 chars on Win2000+; Shell_NotifyIcon truncates safely
    // if the source is shorter (we keep g_tooltip null-terminated).
    strncpy(nid.szTip, g_tooltip, sizeof(nid.szTip) - 1);

    if (!Shell_NotifyIconA(NIM_ADD, &nid)) {
        DestroyWindow(hwnd);
        return 1;
    }
    g_iconLive.store(true, std::memory_order_release);

    MSG msg;
    while (GetMessageA(&msg, NULL, 0, 0) > 0) {
        TranslateMessage(&msg);
        DispatchMessageA(&msg);
    }

    NOTIFYICONDATAA nidDel = {};
    nidDel.cbSize = sizeof(nidDel);
    nidDel.hWnd   = hwnd;
    nidDel.uID    = TRAY_ICON_UID;
    Shell_NotifyIconA(NIM_DELETE, &nidDel);
    g_iconLive.store(false, std::memory_order_release);

    DestroyWindow(hwnd);
    return 0;
}

} // namespace

void StartTrayIcon(const std::string& tooltip, std::function<void()> onExit) {
    if (g_thread) return;  // already started — silently ignore

    g_onExit = std::move(onExit);

    // Snapshot the tooltip into our static buffer so the worker thread
    // doesn't read the caller's std::string.
    strncpy(g_tooltip, tooltip.c_str(), sizeof(g_tooltip) - 1);
    g_tooltip[sizeof(g_tooltip) - 1] = '\0';

    g_thread = CreateThread(NULL, 0, TrayThreadProc, NULL, 0, &g_threadId);
}

void StopTrayIcon() {
    if (!g_thread) return;

    if (g_threadId != 0) {
        // PostThreadMessage delivers a WM_QUIT directly to the tray
        // thread's queue; GetMessage returns 0 and the loop exits.
        PostThreadMessageA(g_threadId, WM_QUIT, 0, 0);
    }
    // 2 s is plenty for the icon-removal + window-destroy path.
    WaitForSingleObject(g_thread, 2000);
    CloseHandle(g_thread);
    g_thread   = nullptr;
    g_threadId = 0;
    g_onExit   = nullptr;  // drop any captured references
}

} // namespace sherpa

#else  // !_WIN32

namespace sherpa {
void StartTrayIcon(const std::string&, std::function<void()>) {}
void StopTrayIcon() {}
} // namespace sherpa

#endif