LenovoLegionToolkit

1
using System;
2
using System.Drawing;
3
using System.Runtime.InteropServices;
4
using System.Threading;
5
using System.Threading.Tasks;
6
using System.Windows;
7
using System.Windows.Controls;
8
using System.Windows.Controls.Primitives;
9
using System.Windows.Forms;
10
using System.Windows.Interop;
11
using LenovoLegionToolkit.Lib.Utils;
12
using Windows.Win32;
13
using Windows.Win32.Foundation;
14
using Windows.Win32.UI.Shell;
15
using Windows.Win32.UI.WindowsAndMessaging;
16
using Wpf.Ui.Controls;
17

18
namespace LenovoLegionToolkit.WPF.Utils;
19

20
public class NotifyIcon : NativeWindow, IDisposable
21
{
22
    private const uint TRAY_MESSAGE_ID = PInvoke.WM_USER + 1069;
23

24
    private static readonly uint TaskbarCreatedMessage = PInvoke.RegisterWindowMessage("TaskbarCreated");
25

26
    private static uint _nextId;
27

28
    private readonly object _lock = new();
29
    private readonly uint _id = ++_nextId;
30

31
    private bool _added;
32
    private CancellationTokenSource? _showToolTipCancellationTokenSource;
33

34
    private bool _visible;
35
    public bool Visible
36
    {
37
        set
38
        {
39
            _visible = value;
40
            UpdateIcon();
41
        }
42
    }
43

44
    private Icon? _icon;
45
    public Icon? Icon
46
    {
47
        set
48
        {
49
            _icon = value;
50
            UpdateIcon();
51

52
        }
53
    }
54

55
    private string? _text;
56
    public string? Text
57
    {
58
        set
59
        {
60
            _text = value;
61
            UpdateIcon();
62
        }
63
    }
64

65
    private UiWindow? _currentToolTipWindow;
66

67
    private Func<Task<UiWindow>>? _toolTipWindow;
68
    public Func<Task<UiWindow>>? ToolTipWindow
69
    {
70
        set
71
        {
72
            _toolTipWindow = value;
73
            UpdateIcon();
74
        }
75
    }
76

77
    public ContextMenu? ContextMenu { get; set; }
78

79
    public event EventHandler? OnClick;
80

81
    public NotifyIcon()
82
    {
83
        UpdateIcon();
84
    }
85

86
    protected override void WndProc(ref Message m)
87
    {
88
        switch ((uint)m.Msg)
89
        {
90
            case TRAY_MESSAGE_ID:
91
                switch ((uint)m.LParam & 0xFFFF)
92
                {
93
                    case PInvoke.NIN_POPUPOPEN:
94
                        if (Log.Instance.IsTraceEnabled)
95
                            Log.Instance.Trace($"NIN_POPUPOPEN");
96
                        ShowToolTip();
97
                        break;
98
                    case PInvoke.NIN_POPUPCLOSE:
99
                        if (Log.Instance.IsTraceEnabled)
100
                            Log.Instance.Trace($"NIN_POPUPCLOSE");
101
                        HideToolTip();
102
                        break;
103
                    case PInvoke.WM_LBUTTONUP:
104
                        if (Log.Instance.IsTraceEnabled)
105
                            Log.Instance.Trace($"WM_LBUTTONUP");
106
                        HideToolTip();
107
                        HideContextMenu();
108
                        OnClick?.Invoke(this, EventArgs.Empty);
109
                        break;
110
                    case PInvoke.WM_RBUTTONUP:
111
                        if (Log.Instance.IsTraceEnabled)
112
                            Log.Instance.Trace($"WM_RBUTTONUP");
113
                        HideToolTip();
114
                        ShowContextMenu();
115
                        break;
116
                }
117
                break;
118
            case PInvoke.WM_DESTROY:
119
                _visible = false;
120
                UpdateIcon();
121
                break;
122
            default:
123
                if (m.Msg == TaskbarCreatedMessage && _visible)
124
                {
125
                    _visible = true;
126
                    _added = false;
127
                    UpdateIcon();
128
                }
129

130
                DefWndProc(ref m);
131
                break;
132
        }
133
    }
134

135
    private async void ShowToolTip()
136
    {
137
        if (_toolTipWindow is null)
138
            return;
139

140
        _showToolTipCancellationTokenSource?.Cancel();
141
        _showToolTipCancellationTokenSource = new();
142

143
        var token = _showToolTipCancellationTokenSource.Token;
144

145
        try
146
        {
147
            await Task.Delay(TimeSpan.FromMilliseconds(500), token);
148

149
            if (ContextMenu is not null && ContextMenu.IsOpen)
150
                return;
151

152
            _currentToolTipWindow?.Close();
153
            _currentToolTipWindow = await _toolTipWindow();
154

155
            token.ThrowIfCancellationRequested();
156

157
            _currentToolTipWindow?.Show();
158
        }
159
        catch (OperationCanceledException)
160
        {
161
            _currentToolTipWindow?.Close();
162
            _currentToolTipWindow = null;
163
        }
164
        catch (Exception ex)
165
        {
166
            _currentToolTipWindow?.Close();
167
            _currentToolTipWindow = null;
168

169
            if (Log.Instance.IsTraceEnabled)
170
                Log.Instance.Trace($"Failed to show tooltip.", ex);
171
        }
172
    }
173

174
    private void HideToolTip()
175
    {
176
        if (_toolTipWindow is null)
177
            return;
178

179
        _showToolTipCancellationTokenSource?.Cancel();
180

181
        _currentToolTipWindow?.Hide();
182
        _currentToolTipWindow = null;
183
    }
184

185
    private void ShowContextMenu()
186
    {
187
        if (ContextMenu is null)
188
            return;
189

190
        ContextMenu.Placement = PlacementMode.Mouse;
191
        ContextMenu.PlacementRectangle = Rect.Empty;
192
        ContextMenu.PlacementTarget = null;
193
        ContextMenu.IsOpen = true;
194

195
        if (PresentationSource.FromVisual(ContextMenu) is HwndSource source && source.Handle != IntPtr.Zero)
196
            PInvoke.SetForegroundWindow(new HWND(source.Handle));
197
    }
198

199
    private void HideContextMenu()
200
    {
201
        if (ContextMenu is null || !ContextMenu.IsOpen)
202
            return;
203

204
        ContextMenu.IsOpen = false;
205
    }
206

207
    private void UpdateIcon()
208
    {
209
        lock (_lock)
210
        {
211
            var data = new NOTIFYICONDATAW
212
            {
213
                cbSize = (uint)Marshal.SizeOf<NOTIFYICONDATAW>(),
214
                uID = _id,
215
                uCallbackMessage = TRAY_MESSAGE_ID,
216
                uFlags = NOTIFY_ICON_DATA_FLAGS.NIF_MESSAGE | NOTIFY_ICON_DATA_FLAGS.NIF_TIP,
217
                szTip = " "
218
            };
219

220
            if (_visible && Handle == IntPtr.Zero)
221
                CreateHandle(new CreateParams());
222

223
            data.hWnd = new HWND(Handle);
224

225
            if (_icon is not null)
226
            {
227
                data.uFlags |= NOTIFY_ICON_DATA_FLAGS.NIF_ICON;
228
                data.hIcon = new HICON(_icon.Handle);
229
            }
230

231
            if (_text is not null && _toolTipWindow is null)
232
            {
233
                data.uFlags |= NOTIFY_ICON_DATA_FLAGS.NIF_SHOWTIP;
234
                data.szTip = _text;
235
            }
236

237
            switch (_visible, _added)
238
            {
239
                case (true, false):
240
                    PInvoke.Shell_NotifyIcon(NOTIFY_ICON_MESSAGE.NIM_ADD, data);
241
                    data.Anonymous = new() { uVersion = 4 };
242
                    PInvoke.Shell_NotifyIcon(NOTIFY_ICON_MESSAGE.NIM_SETVERSION, data);
243
                    _added = true;
244
                    break;
245
                case (true, true):
246
                    PInvoke.Shell_NotifyIcon(NOTIFY_ICON_MESSAGE.NIM_MODIFY, data);
247
                    break;
248
                case (false, true):
249
                    PInvoke.Shell_NotifyIcon(NOTIFY_ICON_MESSAGE.NIM_DELETE, data);
250
                    _added = false;
251
                    break;
252
            }
253
        }
254
    }
255

256
    public void Dispose()
257
    {
258
        GC.SuppressFinalize(this);
259

260
        HideContextMenu();
261
        HideToolTip();
262

263
        _visible = false;
264
        UpdateIcon();
265

266
        _icon?.Dispose();
267

268
        _icon = null;
269
        _text = null;
270
        _toolTipWindow = null;
271
        ContextMenu = null;
272

273
        ReleaseHandle();
274
    }
275
}
276