// Pin a Window (DeskTools module): keep any window on top of all the others.
// A small blue box with a plus sign sits on the title bar of every pinned window.
// Click the box to unpin. Ctrl+Alt+P pins or unpins the active window.

#ifndef UNICODE
#define UNICODE
#endif
#ifndef _UNICODE
#define _UNICODE
#endif
#define WIN32_LEAN_AND_MEAN
#define _WIN32_WINNT 0x0A00
#include <windows.h>
#include "desktools.h"
#include <dwmapi.h>
#include <stdio.h>
#include <wchar.h>

#define WM_PINAT      (WM_APP + 1)
#define WM_UNPINIDX   (WM_APP + 2)
#define WM_REPOS      (WM_APP + 3)
#define MAX_PINS      64
#define HK_TOGGLE     1
#define HK_CANCEL     2
#define TIMER_CHECK   1
#define TIMER_OSD     2
#define TIMER_FOLLOW  3

#define IDM_PICK      100
#define IDM_UNPINALL  101
#define IDM_PIN_BASE  1000

typedef struct {
    BOOL used;
    HWND target;
    HWND badge;
    BOOL wasTopmost;
    BOOL hover;
} Pin;

static Pin       g_pins[MAX_PINS];
static HINSTANCE g_inst;
static HWND      g_main, g_osd;
static BOOL      g_enabled = TRUE, g_hotkeyOk, g_picking, g_swallowUp;
static HWINEVENTHOOK g_evFg, g_evLoc, g_evMin;
static int       g_dpi = 96;
static HFONT     g_osdFont;
static WCHAR     g_osdText[128];

static int S(int v) { return MulDiv(v, g_dpi, 96); }
static int BadgeSize(void) { return S(18); }

/* ---------------- small on-screen message ---------------- */

static void ShowOsdAt(int cx, int cy, const WCHAR *text, UINT ms) {
    wcsncpy(g_osdText, text, 127); g_osdText[127] = 0;
    HDC dc = GetDC(NULL);
    SelectObject(dc, g_osdFont);
    SIZE sz; GetTextExtentPoint32W(dc, g_osdText, (int)wcslen(g_osdText), &sz);
    ReleaseDC(NULL, dc);
    int w = sz.cx + S(28), h = S(40);
    HRGN rgn = CreateRoundRectRgn(0, 0, w + 1, h + 1, S(12), S(12));
    SetWindowRgn(g_osd, rgn, FALSE);
    SetWindowPos(g_osd, HWND_TOPMOST, cx - w / 2, cy - h / 2, w, h, SWP_NOACTIVATE | SWP_SHOWWINDOW);
    InvalidateRect(g_osd, NULL, TRUE);
    UpdateWindow(g_osd);
    if (ms) SetTimer(g_main, TIMER_OSD, ms, NULL); else KillTimer(g_main, TIMER_OSD);
}

static void ShowOsdOver(HWND target, const WCHAR *text) {
    RECT r;
    if (target && GetWindowRect(target, &r)) ShowOsdAt((r.left + r.right) / 2, (r.top + r.bottom) / 2, text, 1100);
    else { POINT p; GetCursorPos(&p); ShowOsdAt(p.x, p.y + S(40), text, 1100); }
}

static LRESULT CALLBACK OsdProc(HWND h, UINT msg, WPARAM wp, LPARAM lp) {
    switch (msg) {
    case WM_PAINT: {
        PAINTSTRUCT ps; HDC dc = BeginPaint(h, &ps);
        RECT rc; GetClientRect(h, &rc);
        HBRUSH bg = CreateSolidBrush(RGB(32, 36, 44));
        FillRect(dc, &rc, bg); DeleteObject(bg);
        SelectObject(dc, g_osdFont);
        SetBkMode(dc, TRANSPARENT);
        SetTextColor(dc, RGB(255, 255, 255));
        DrawTextW(dc, g_osdText, -1, &rc, DT_CENTER | DT_VCENTER | DT_SINGLELINE | DT_NOPREFIX);
        EndPaint(h, &ps);
        return 0;
    }
    case WM_MOUSEACTIVATE: return MA_NOACTIVATE;
    case WM_NCHITTEST: return HTTRANSPARENT;
    }
    return DefWindowProcW(h, msg, wp, lp);
}

/* ---------------- the box-with-a-plus badge ---------------- */

static void RenderBadge(Pin *p) {
    int B = BadgeSize();
    HDC screen = GetDC(NULL);
    BITMAPINFO bi; ZeroMemory(&bi, sizeof(bi));
    bi.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
    bi.bmiHeader.biWidth = B; bi.bmiHeader.biHeight = -B;
    bi.bmiHeader.biPlanes = 1; bi.bmiHeader.biBitCount = 32; bi.bmiHeader.biCompression = BI_RGB;
    BYTE *px = NULL;
    HDC dc = CreateCompatibleDC(screen);
    HBITMAP bmp = CreateDIBSection(screen, &bi, DIB_RGB_COLORS, (void **)&px, NULL, 0);
    if (bmp) {
        HGDIOBJ old = SelectObject(dc, bmp);
        // box: rounded square, blue fill, white edge; plus sign: white bars
        int r = S(4), edge = S(1) > 1 ? S(1) : 1;
        int bar = S(2) > 2 ? S(2) : 2, len = B * 5 / 9;
        int c = B / 2, b0 = c - bar / 2, b1 = b0 + bar, l0 = c - len / 2, l1 = l0 + len;
        int fr = p->hover ? 70 : 40, fg = p->hover ? 145 : 115, fb = p->hover ? 245 : 215;
        for (int y = 0; y < B; y++) {
            for (int x = 0; x < B; x++) {
                int i = (y * B + x) * 4;
                // distance outside the rounded corner
                int dx = x < r ? r - x : (x >= B - r ? x - (B - r - 1) : 0);
                int dy = y < r ? r - y : (y >= B - r ? y - (B - r - 1) : 0);
                int inside = (dx * dx + dy * dy) <= r * r;
                int a = 0, R = 0, G = 0, Bl = 0;
                if (inside) {
                    int onEdge = x < edge || y < edge || x >= B - edge || y >= B - edge ||
                                 (dx && dy && (dx * dx + dy * dy) > (r - edge) * (r - edge));
                    int plus = (x >= b0 && x < b1 && y >= l0 && y < l1) || (y >= b0 && y < b1 && x >= l0 && x < l1);
                    if (onEdge || plus) { R = G = Bl = 255; a = 255; }
                    else { R = fr; G = fg; Bl = fb; a = 235; }
                }
                px[i] = (BYTE)(Bl * a / 255); px[i + 1] = (BYTE)(G * a / 255); px[i + 2] = (BYTE)(R * a / 255); px[i + 3] = (BYTE)a;
            }
        }
        SIZE sz = { B, B }; POINT src = { 0, 0 };
        BLENDFUNCTION bf = { AC_SRC_OVER, 0, 255, AC_SRC_ALPHA };
        UpdateLayeredWindow(p->badge, screen, NULL, &sz, dc, &src, 0, &bf, ULW_ALPHA);
        SelectObject(dc, old);
        DeleteObject(bmp);
    }
    DeleteDC(dc);
    ReleaseDC(NULL, screen);
}

static void PlaceBadge(Pin *p) {
    if (!p->used || !p->badge) return;
    HWND t = p->target;
    if (!IsWindow(t)) return;
    if (IsIconic(t) || !IsWindowVisible(t) || !g_enabled) {
        if (IsWindowVisible(p->badge)) ShowWindow(p->badge, SW_HIDE);
        return;
    }
    int B = BadgeSize();
    RECT wr; GetWindowRect(t, &wr);
    RECT fr = wr;
    DwmGetWindowAttribute(t, DWMWA_EXTENDED_FRAME_BOUNDS, &fr, sizeof(fr));
    RECT cb; ZeroMemory(&cb, sizeof(cb));
    int x, y;
    if (SUCCEEDED(DwmGetWindowAttribute(t, DWMWA_CAPTION_BUTTON_BOUNDS, &cb, sizeof(cb))) && cb.right > cb.left) {
        // just to the left of the minimize / maximize / close buttons
        x = wr.left + cb.left - B - S(8);
        y = wr.top + cb.top + ((cb.bottom - cb.top) - B) / 2;
    } else {
        x = fr.right - S(150) - B;
        y = fr.top + S(7);
    }
    if (x < fr.left + S(4)) x = fr.left + S(4);
    if (y < fr.top + S(2)) y = fr.top + S(2);
    SetWindowPos(p->badge, HWND_TOPMOST, x, y, 0, 0, SWP_NOSIZE | SWP_NOACTIVATE | SWP_SHOWWINDOW);
}

static Pin *PinOfBadge(HWND b) {
    for (int i = 0; i < MAX_PINS; i++) if (g_pins[i].used && g_pins[i].badge == b) return &g_pins[i];
    return NULL;
}

static LRESULT CALLBACK BadgeProc(HWND h, UINT msg, WPARAM wp, LPARAM lp) {
    Pin *p = PinOfBadge(h);
    switch (msg) {
    case WM_MOUSEACTIVATE: return MA_NOACTIVATE;
    case WM_SETCURSOR: SetCursor(LoadCursorW(NULL, IDC_HAND)); return TRUE;
    case WM_MOUSEMOVE:
        if (p && !p->hover) {
            p->hover = TRUE;
            TRACKMOUSEEVENT t = { sizeof(t), TME_LEAVE, h, 0 };
            TrackMouseEvent(&t);
            RenderBadge(p);
        }
        return 0;
    case WM_MOUSELEAVE:
        if (p && p->hover) { p->hover = FALSE; RenderBadge(p); }
        return 0;
    case WM_LBUTTONUP:
        if (p) PostMessageW(g_main, WM_UNPINIDX, (WPARAM)(p - g_pins), 0);
        return 0;
    }
    return DefWindowProcW(h, msg, wp, lp);
}

/* ---------------- pin / unpin ---------------- */

static Pin *FindPin(HWND t) {
    for (int i = 0; i < MAX_PINS; i++) if (g_pins[i].used && g_pins[i].target == t) return &g_pins[i];
    return NULL;
}

static int CountPins(void) {
    int n = 0;
    for (int i = 0; i < MAX_PINS; i++) if (g_pins[i].used) n++;
    return n;
}

static BOOL Pinnable(HWND h) {
    if (!h || !IsWindow(h) || !IsWindowVisible(h)) return FALSE;
    DWORD pid = 0; GetWindowThreadProcessId(h, &pid);
    if (pid == GetCurrentProcessId()) return FALSE;
    WCHAR cls[64]; GetClassNameW(h, cls, 64);
    if (!wcscmp(cls, L"Progman") || !wcscmp(cls, L"WorkerW") || !wcscmp(cls, L"Shell_TrayWnd") ||
        !wcscmp(cls, L"Shell_SecondaryTrayWnd"))
        return FALSE;
    return TRUE;
}

static void Unpin(Pin *p, BOOL quiet) {
    if (!p->used) return;
    HWND t = p->target;
    if (IsWindow(t) && !p->wasTopmost)
        SetWindowPos(t, HWND_NOTOPMOST, 0, 0, 0, 0, SWP_NOMOVE | SWP_NOSIZE | SWP_NOACTIVATE);
    if (p->badge) DestroyWindow(p->badge);
    p->used = FALSE; p->badge = NULL;
    if (!quiet && IsWindow(t)) ShowOsdOver(t, L"Unpinned");
}

static void PinWindow(HWND h) {
    h = GetAncestor(h, GA_ROOT);
    if (!Pinnable(h)) return;
    Pin *existing = FindPin(h);
    if (existing) { Unpin(existing, FALSE); return; }
    Pin *p = NULL;
    for (int i = 0; i < MAX_PINS; i++) if (!g_pins[i].used) { p = &g_pins[i]; break; }
    if (!p) return;
    ZeroMemory(p, sizeof(*p));
    p->target = h;
    p->wasTopmost = (GetWindowLongW(h, GWL_EXSTYLE) & WS_EX_TOPMOST) != 0;
    SetWindowPos(h, HWND_TOPMOST, 0, 0, 0, 0, SWP_NOMOVE | SWP_NOSIZE | SWP_NOACTIVATE);
    if (!(GetWindowLongW(h, GWL_EXSTYLE) & WS_EX_TOPMOST)) {
        ShowOsdOver(h, L"Can't pin this window (admin program?)");
        return;
    }
    p->badge = CreateWindowExW(WS_EX_LAYERED | WS_EX_TOOLWINDOW | WS_EX_TOPMOST | WS_EX_NOACTIVATE,
                               L"DeskToolsPinBadge", L"", WS_POPUP, 0, 0, BadgeSize(), BadgeSize(),
                               NULL, NULL, g_inst, NULL);
    p->used = TRUE;
    RenderBadge(p);
    PlaceBadge(p);
    ShowOsdOver(h, L"Pinned on top");
}

static void UnpinAll(BOOL quiet) {
    for (int i = 0; i < MAX_PINS; i++) if (g_pins[i].used) Unpin(&g_pins[i], quiet);
}

/* ---------------- picking a window with the mouse ---------------- */

static void StopPicking(void) {
    if (!g_picking) return;
    g_picking = FALSE;
    KillTimer(g_main, TIMER_FOLLOW);
    UnregisterHotKey(g_main, HK_CANCEL);
    ShowWindow(g_osd, SW_HIDE);
}

static void StartPicking(void) {
    if (!g_enabled) return;
    g_picking = TRUE;
    RegisterHotKey(g_main, HK_CANCEL, 0, VK_ESCAPE);
    SetTimer(g_main, TIMER_FOLLOW, 30, NULL);
    POINT pt; GetCursorPos(&pt);
    ShowOsdAt(pt.x, pt.y + S(44), L"Click a window to pin it on top  (Esc cancels)", 0);
}

int PW_Mouse(WPARAM wp, MSLLHOOKSTRUCT *m) {
    if (!g_picking) {
        if (g_swallowUp && wp == WM_LBUTTONUP) { g_swallowUp = FALSE; return 1; }
        return 0;
    }
    if (wp == WM_LBUTTONDOWN) {
        PostMessageW(g_main, WM_PINAT, 0, MAKELPARAM((WORD)(SHORT)m->pt.x, (WORD)(SHORT)m->pt.y));
        g_swallowUp = TRUE;
        return 1;
    }
    if (wp == WM_RBUTTONDOWN || wp == WM_RBUTTONUP) {
        PostMessageW(g_main, WM_PINAT, 1, 0);   // right-click cancels
        return 1;
    }
    return 0;
}

/* ---------------- keeping badges in place ---------------- */

static void CALLBACK WinEventProc(HWINEVENTHOOK hook, DWORD ev, HWND h, LONG obj, LONG child, DWORD tid, DWORD time) {
    (void)hook; (void)child; (void)tid; (void)time;
    if (ev == EVENT_SYSTEM_FOREGROUND) {
        // a pinned window that comes to the front would cover its badge, so lift the badges back up
        PostMessageW(g_main, WM_REPOS, 0, 0);
        return;
    }
    if (obj != OBJID_WINDOW || !h) return;
    Pin *p = FindPin(h);
    if (p) PlaceBadge(p);
}

static LRESULT CALLBACK MainProc(HWND h, UINT msg, WPARAM wp, LPARAM lp) {
    switch (msg) {
    case WM_HOTKEY:
        if (wp == HK_TOGGLE) {
            HWND fg = GetForegroundWindow();
            if (fg) PinWindow(fg);
        } else if (wp == HK_CANCEL) {
            StopPicking();
        }
        return 0;
    case WM_PINAT:
        StopPicking();
        if (wp == 0) {
            POINT pt = { (short)LOWORD(lp), (short)HIWORD(lp) };
            HWND w = WindowFromPoint(pt);
            if (w) PinWindow(w);
        }
        return 0;
    case WM_UNPINIDX:
        if (wp < MAX_PINS) Unpin(&g_pins[wp], FALSE);
        return 0;
    case WM_REPOS:
        for (int i = 0; i < MAX_PINS; i++) if (g_pins[i].used) PlaceBadge(&g_pins[i]);
        return 0;
    case WM_TIMER:
        if (wp == TIMER_OSD) { KillTimer(h, TIMER_OSD); if (!g_picking) ShowWindow(g_osd, SW_HIDE); }
        else if (wp == TIMER_FOLLOW && g_picking) {
            POINT pt; GetCursorPos(&pt);
            RECT r; GetWindowRect(g_osd, &r);
            SetWindowPos(g_osd, HWND_TOPMOST, pt.x - (r.right - r.left) / 2, pt.y + S(24), 0, 0,
                         SWP_NOSIZE | SWP_NOACTIVATE);
        } else if (wp == TIMER_CHECK) {
            for (int i = 0; i < MAX_PINS; i++) {
                Pin *p = &g_pins[i];
                if (!p->used) continue;
                if (!IsWindow(p->target)) { if (p->badge) DestroyWindow(p->badge); p->used = FALSE; continue; }
                // something else may have dropped the window off the top; put it back
                if (!(GetWindowLongW(p->target, GWL_EXSTYLE) & WS_EX_TOPMOST))
                    SetWindowPos(p->target, HWND_TOPMOST, 0, 0, 0, 0, SWP_NOMOVE | SWP_NOSIZE | SWP_NOACTIVATE);
                PlaceBadge(p);
            }
        }
        return 0;
    case WM_COMMAND: {
        UINT id = LOWORD(wp);
        if (id >= IDM_PIN_BASE && id < IDM_PIN_BASE + MAX_PINS) { Unpin(&g_pins[id - IDM_PIN_BASE], FALSE); return 0; }
        if (id == IDM_PICK) StartPicking();
        else if (id == IDM_UNPINALL) UnpinAll(TRUE);
        return 0;
    }
    case WM_QUERYENDSESSION:
        UnpinAll(TRUE);
        return TRUE;
    case WM_DESTROY:
        UnpinAll(TRUE);
        return 0;
    }
    return DefWindowProcW(h, msg, wp, lp);
}

/* ---------------- DeskTools module interface ---------------- */

void PW_SetEnabled(BOOL on) {
    g_enabled = on;
    if (on) {
        if (!g_hotkeyOk) g_hotkeyOk = RegisterHotKey(g_main, HK_TOGGLE, MOD_CONTROL | MOD_ALT | MOD_NOREPEAT, 'P');
    } else {
        StopPicking();
        UnpinAll(TRUE);
        if (g_hotkeyOk) { UnregisterHotKey(g_main, HK_TOGGLE); g_hotkeyOk = FALSE; }
    }
}
BOOL PW_IsEnabled(void) { return g_enabled; }

void PW_Init(HINSTANCE inst, BOOL enabled) {
    g_inst = inst;
    HDC sdc = GetDC(NULL); g_dpi = GetDeviceCaps(sdc, LOGPIXELSY); ReleaseDC(NULL, sdc);
    if (g_dpi < 96) g_dpi = 96;
    g_osdFont = CreateFontW(-MulDiv(11, g_dpi, 72), 0, 0, 0, FW_SEMIBOLD, FALSE, FALSE, FALSE, DEFAULT_CHARSET,
                            OUT_DEFAULT_PRECIS, CLIP_DEFAULT_PRECIS, CLEARTYPE_QUALITY, DEFAULT_PITCH, L"Segoe UI");

    WNDCLASSEXW wc; ZeroMemory(&wc, sizeof(wc));
    wc.cbSize = sizeof(wc); wc.lpfnWndProc = MainProc; wc.hInstance = inst; wc.lpszClassName = L"DeskToolsPinMain";
    RegisterClassExW(&wc);
    WNDCLASSEXW bc; ZeroMemory(&bc, sizeof(bc));
    bc.cbSize = sizeof(bc); bc.lpfnWndProc = BadgeProc; bc.hInstance = inst; bc.lpszClassName = L"DeskToolsPinBadge";
    bc.hCursor = LoadCursorW(NULL, IDC_HAND);
    RegisterClassExW(&bc);
    WNDCLASSEXW oc; ZeroMemory(&oc, sizeof(oc));
    oc.cbSize = sizeof(oc); oc.lpfnWndProc = OsdProc; oc.hInstance = inst; oc.lpszClassName = L"DeskToolsPinOSD";
    RegisterClassExW(&oc);

    g_main = CreateWindowExW(0, L"DeskToolsPinMain", L"Pin a Window", 0, 0, 0, 0, 0, NULL, NULL, inst, NULL);
    g_osd = CreateWindowExW(WS_EX_TOPMOST | WS_EX_TOOLWINDOW | WS_EX_NOACTIVATE | WS_EX_LAYERED | WS_EX_TRANSPARENT,
                            L"DeskToolsPinOSD", L"", WS_POPUP, 0, 0, 10, 10, NULL, NULL, inst, NULL);
    SetLayeredWindowAttributes(g_osd, 0, 230, LWA_ALPHA);

    g_evFg  = SetWinEventHook(EVENT_SYSTEM_FOREGROUND, EVENT_SYSTEM_FOREGROUND, NULL, WinEventProc, 0, 0, WINEVENT_OUTOFCONTEXT);
    g_evLoc = SetWinEventHook(EVENT_OBJECT_LOCATIONCHANGE, EVENT_OBJECT_LOCATIONCHANGE, NULL, WinEventProc, 0, 0, WINEVENT_OUTOFCONTEXT);
    g_evMin = SetWinEventHook(EVENT_SYSTEM_MINIMIZESTART, EVENT_SYSTEM_MINIMIZEEND, NULL, WinEventProc, 0, 0, WINEVENT_OUTOFCONTEXT);
    SetTimer(g_main, TIMER_CHECK, 500, NULL);
    g_enabled = !enabled;
    PW_SetEnabled(enabled);
}

void PW_Shutdown(void) {
    if (g_evFg) UnhookWinEvent(g_evFg);
    if (g_evLoc) UnhookWinEvent(g_evLoc);
    if (g_evMin) UnhookWinEvent(g_evMin);
    g_evFg = g_evLoc = g_evMin = NULL;
    if (g_main) { DestroyWindow(g_main); g_main = NULL; }
    if (g_osd) { DestroyWindow(g_osd); g_osd = NULL; }
    if (g_osdFont) { DeleteObject(g_osdFont); g_osdFont = NULL; }
}

void PW_AppendMenu(HMENU parent, UINT base) {
    HMENU m = CreatePopupMenu();
    UINT dis = g_enabled ? 0 : MF_GRAYED;
    AppendMenuW(m, MF_STRING | (g_enabled ? MF_CHECKED : 0), base + DT_ID_ENABLE, L"Turned on");
    AppendMenuW(m, MF_SEPARATOR, 0, NULL);
    AppendMenuW(m, MF_STRING | dis, base + IDM_PICK, L"Pin a window...\tCtrl+Alt+P");
    int n = 0;
    for (int i = 0; i < MAX_PINS; i++) {
        Pin *p = &g_pins[i];
        if (!p->used || !IsWindow(p->target)) continue;
        WCHAR title[80], item[128];
        GetWindowTextW(p->target, title, 60);
        if (!title[0]) wcscpy(title, L"(untitled)");
        swprintf(item, 128, L"Unpin: %ls", title);
        AppendMenuW(m, MF_STRING, base + IDM_PIN_BASE + i, item);
        n++;
    }
    AppendMenuW(m, MF_STRING | (n ? 0 : MF_GRAYED), base + IDM_UNPINALL, L"Unpin all windows");
    WCHAR label[64];
    if (n) swprintf(label, 64, L"Pin a Window  (%d pinned)", n); else wcscpy(label, L"Pin a Window");
    AppendMenuW(parent, MF_POPUP, (UINT_PTR)m, label);
    (void)CountPins;
}

void PW_Command(UINT local) {
    if (local == DT_ID_ENABLE) { PW_SetEnabled(!g_enabled); return; }
    SendMessageW(g_main, WM_COMMAND, local, 0);
}