main.cpp

#include <windows.h>
#include <pdh.h>
#include <thread>
#include <chrono>
#include <cstdio>
#include <sstream>

#pragma comment(lib, "pdh.lib")
#pragma comment(linker, "/SUBSYSTEM:WINDOWS /ENTRY:WinMainCRTStartup")

HINSTANCE hInst;
NOTIFYICONDATAA nid_cpu;
NOTIFYICONDATAA nid_gpu;
PDH_HQUERY cpuQuery;
PDH_HCOUNTER cpuTotal;

HMENU hTrayMenu = NULL;

float GetCPULoad() {
    FILETIME idleTime, kernelTime, userTime;
    static FILETIME prevIdleTime = {}, prevKernelTime = {}, prevUserTime = {};
    static bool first = true;

    if (first) {
        GetSystemTimes(&prevIdleTime, &prevKernelTime, &prevUserTime);
        first = false;
        Sleep(1000);
        GetSystemTimes(&idleTime, &kernelTime, &userTime);
    } else {
        GetSystemTimes(&idleTime, &kernelTime, &userTime);
    }

    auto FileTimeToULL = [](const FILETIME& ft) {
        return (((ULONGLONG)ft.dwHighDateTime) << 32) | ft.dwLowDateTime;
    };

    ULONGLONG idle = FileTimeToULL(idleTime) - FileTimeToULL(prevIdleTime);
    ULONGLONG kernel = FileTimeToULL(kernelTime) - FileTimeToULL(prevKernelTime);
    ULONGLONG user = FileTimeToULL(userTime) - FileTimeToULL(prevUserTime);

    ULONGLONG total = kernel + user;

    float cpuUsage = 0.0f;
    if (total != 0)
        cpuUsage = (float)(total - idle) * 100.0f / total;

    prevIdleTime = idleTime;
    prevKernelTime = kernelTime;
    prevUserTime = userTime;

    return cpuUsage;
}

float GetGPULoad() {
    char cmd[] = "nvidia-smi --query-gpu=utilization.gpu --format=csv,noheader,nounits";
    char buffer[128] = {0};
    SECURITY_ATTRIBUTES sa = { sizeof(SECURITY_ATTRIBUTES), NULL, TRUE };
    HANDLE hRead = NULL, hWrite = NULL;
    if (!CreatePipe(&hRead, &hWrite, &sa, 0)) return -1.0f;

    STARTUPINFOA si = {};
    si.cb = sizeof(si);
    si.dwFlags = STARTF_USESHOWWINDOW | STARTF_USESTDHANDLES;
    si.wShowWindow = SW_HIDE;
    si.hStdOutput = hWrite;
    si.hStdError = hWrite;

    PROCESS_INFORMATION pi = {};
    BOOL ok = CreateProcessA(
        NULL, cmd, NULL, NULL, TRUE,
        CREATE_NO_WINDOW, NULL, NULL, &si, &pi
    );
    if (!ok) {
        CloseHandle(hWrite);
        CloseHandle(hRead);
        return -1.0f;
    }

    CloseHandle(hWrite);

    DWORD read = 0;
    ReadFile(hRead, buffer, sizeof(buffer) - 1, &read, NULL);
    buffer[read] = 0;

    CloseHandle(hRead);
    WaitForSingleObject(pi.hProcess, 500);
    CloseHandle(pi.hProcess);
    CloseHandle(pi.hThread);

    int gpuUtil = -1;
    std::istringstream iss(buffer);
    iss >> gpuUtil;
    return (float)gpuUtil;
}

HICON CreateValueIcon(int value, const char* label) {
    const int size = 16;
    HDC hdc = CreateCompatibleDC(NULL);
    HBITMAP hBmp = CreateCompatibleBitmap(GetDC(NULL), size, size);
    SelectObject(hdc, hBmp);

    HBRUSH brush = CreateSolidBrush(RGB(0,0,0));
    RECT r = {0,0,size,size};
    FillRect(hdc, &r, brush);
    DeleteObject(brush);

    SetTextColor(hdc, RGB(0,255,0));
    SetBkMode(hdc, TRANSPARENT);
    char buf[8];
    sprintf(buf, "%s%d", label, value);
    DrawTextA(hdc, buf, -1, &r, DT_CENTER | DT_VCENTER | DT_SINGLELINE);

    ICONINFO ii = {};
    ii.fIcon = TRUE;
    ii.hbmColor = hBmp;
    ii.hbmMask = hBmp;
    HICON hIcon = CreateIconIndirect(&ii);

    DeleteObject(hBmp);
    DeleteDC(hdc);

    return hIcon;
}

void CreateTrayIcons(HWND hwnd) {
    ZeroMemory(&nid_cpu, sizeof(nid_cpu));
    nid_cpu.cbSize = sizeof(NOTIFYICONDATAA);
    nid_cpu.hWnd = hwnd;
    nid_cpu.uID = 1;
    nid_cpu.uFlags = NIF_ICON | NIF_TIP | NIF_MESSAGE;
    nid_cpu.uCallbackMessage = WM_USER + 1;
    nid_cpu.hIcon = CreateValueIcon(0, "C");
    strcpy(nid_cpu.szTip, "CPU Monitor");
    Shell_NotifyIconA(NIM_ADD, &nid_cpu);

    ZeroMemory(&nid_gpu, sizeof(nid_gpu));
    nid_gpu.cbSize = sizeof(NOTIFYICONDATAA);
    nid_gpu.hWnd = hwnd;
    nid_gpu.uID = 2;
    nid_gpu.uFlags = NIF_ICON | NIF_TIP | NIF_MESSAGE;
    nid_gpu.uCallbackMessage = WM_USER + 2;
    nid_gpu.hIcon = CreateValueIcon(0, "G");
    strcpy(nid_gpu.szTip, "GPU Monitor");
    Shell_NotifyIconA(NIM_ADD, &nid_gpu);
}

void UpdateTrayIcons(int cpuLoad, int gpuLoad) {
    HICON newCpuIcon = CreateValueIcon(cpuLoad, "");
    nid_cpu.hIcon = newCpuIcon;
    Shell_NotifyIconA(NIM_MODIFY, &nid_cpu);
    DestroyIcon(newCpuIcon);

    HICON newGpuIcon = CreateValueIcon(gpuLoad, "");
    nid_gpu.hIcon = newGpuIcon;
    Shell_NotifyIconA(NIM_MODIFY, &nid_gpu);
    DestroyIcon(newGpuIcon);
}

void ShowTrayMenu(HWND hwnd, POINT pt) {
    if (!hTrayMenu) {
        hTrayMenu = CreatePopupMenu();
        AppendMenuA(hTrayMenu, MF_STRING, 1, "Exit");
    }
    SetForegroundWindow(hwnd);
    TrackPopupMenu(hTrayMenu, TPM_RIGHTBUTTON, pt.x, pt.y, 0, hwnd, NULL);
}

LRESULT CALLBACK WndProc(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
    switch (msg) {
    case WM_DESTROY:
        Shell_NotifyIconA(NIM_DELETE, &nid_cpu);
        Shell_NotifyIconA(NIM_DELETE, &nid_gpu);
        if (hTrayMenu) DestroyMenu(hTrayMenu);
        PostQuitMessage(0);
        break;
    case WM_USER + 1: // CPU tray icon callback
    case WM_USER + 2: // GPU tray icon callback
        if (lParam == WM_RBUTTONUP) {
            POINT pt;
            GetCursorPos(&pt);
            ShowTrayMenu(hwnd, pt);
        }
        break;
    case WM_COMMAND:
        if (LOWORD(wParam) == 1) { // Exit
            PostQuitMessage(0);
        }
        break;
    default:
        return DefWindowProc(hwnd, msg, wParam, lParam);
    }
    return 0;
}

int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE, LPSTR, int) {
    hInst = hInstance;

    PdhOpenQueryA(NULL, 0, &cpuQuery);
    PdhAddCounterA(cpuQuery, "\\Processor(_Total)\\% Processor Time", 0, &cpuTotal);
    PdhCollectQueryData(cpuQuery);

    const char CLASS_NAME[] = "CPUTrayClass";
    WNDCLASSA wc = {};
    wc.lpfnWndProc = WndProc;
    wc.hInstance = hInstance;
    wc.lpszClassName = CLASS_NAME;
    RegisterClassA(&wc);

    HWND hwnd = CreateWindowExA(0, CLASS_NAME, "CPU Tray", 0, 0, 0, 0, 0, HWND_MESSAGE, NULL, hInstance, NULL);
    if (!hwnd) return 0;

    CreateTrayIcons(hwnd);

    // Zmena cyklu
    MSG msg;
    while (true) {
        while (PeekMessage(&msg, NULL, 0, 0, PM_REMOVE)) {
            if (msg.message == WM_QUIT)
                return 0;
            TranslateMessage(&msg);
            DispatchMessage(&msg);
        }

        float cpu = GetCPULoad();
        float gpu = GetGPULoad();
        UpdateTrayIcons((int)cpu, (int)gpu);

        std::this_thread::sleep_for(std::chrono::milliseconds(1000));
    }
}