1 files changed, 64 insertions, 7 deletions

diff --git a/drm.ahk b/drm.ahk
index c373b2c..373353f 100644
--- a/drm.ahk
+++ b/drm.ahk
@@ -44,18 +44,46 @@ Menu, Tray, NoStandard
 Menu, Tray, Icon, % "HBITMAP:*" . Create_tray_ico()
 Menu, Tray, Tip, Dynamic Resolution Manager
 for k, res in Resolutions
-  Menu, Tray, Add, %res%, SelectMenuResolution
+  if (res = "")
+    Menu, Tray, Add
+  else
+    Menu, Tray, Add, %res%, ButtonResolution
 Menu, Tray, Add
-Menu, Tray, Add, &Configure..., Configure
+Menu, Tray, Add, &Configure..., ButtonConfigure
 Menu, Tray, Add, E&xit, ButtonExit
+AddCheckmark()
+OnMessage(0x404, "NotifyIcon")
 return
 
 ButtonExit:
 ExitApp
 
+ButtonRestart:
+Reload
+return
+
+; Handle tray opening
+
+NotifyIcon(wParam, lParam)
+{
+  if (lParam = 0x205) ; WM_RBUTTONUP
+  {
+    Menu, Tray, UseErrorLevel
+    Menu, Tray, Delete, &Restart
+    Menu, Tray, Delete, E&xit
+    Menu, Tray, UseErrorLevel, Off
+    if (GetKeyState("Shift") = 1)
+      Menu, Tray, Add, &Restart, ButtonRestart
+    else
+      Menu, Tray, Add, E&xit, ButtonExit
+    Menu, Tray, Show
+    return 1
+  }
+}
+
 ; Open configuration
 
-Configure()
+ButtonConfigure()
 {
   RegWrite, REG_SZ, HKEY_CURRENT_USER\Software\Microsoft\Windows\CurrentVersion\Applets\Regedit, LastKey, HKEY_CURRENT_USER\Software\JohnAJ\DRM
   Run, regedit /m
@@ -63,7 +91,7 @@ Configure()
 
 ; Select resolution from tray menu
 
-SelectMenuResolution(ItemName, ItemPos, MenuName)
+ButtonResolution(ItemName, ItemPos, MenuName)
 {
   w := 0
   h := 0
@@ -86,6 +114,30 @@ SelectMenuResolution(ItemName, ItemPos, MenuName)
   SwitchResolution(new, old)
 }
 
+; Add checkmark to current resolution (if available)
+
+AddCheckmark(res := 0)
+{
+  if (res = 0)
+    res := Get()
+  Menu, Tray, UseErrorLevel
+  Menu, Tray, Check, % res.width "x" res.height " " res.rate "Hz"
+  if (ErrorLevel = 1)
+    Menu, Tray, Check, % res.width "x" res.height
+  Menu, Tray, UseErrorLevel, Off
+}
+
+RemoveCheckmark(res := 0)
+{
+  if (res = 0)
+    res := Get()
+  Menu, Tray, UseErrorLevel
+  Menu, Tray, Uncheck, % res.width "x" res.height " " res.rate "Hz"
+  if (ErrorLevel = 1)
+    Menu, Tray, Uncheck, % res.width "x" res.height
+  Menu, Tray, UseErrorLevel, Off
+}
+
 ; Save window positions, switch to given resolution, restore window positions
 
 SwitchResolution(new, old)
@@ -96,6 +148,8 @@ SwitchResolution(new, old)
   Set(new)
   Sleep, 500
   Restore(Positions)
+  RemoveCheckmark(old)
+  AddCheckmark(new)
   
   TrayTip, Resolution changed, % "Press Escape within 3 seconds to revert."
   KeyWait, Esc, D T3
@@ -106,6 +160,8 @@ SwitchResolution(new, old)
   Set(old)
   Sleep, 500
   Restore(Positions)
+  RemoveCheckmark(new)
+  AddCheckmark(old)
 }
 
 Save(ByRef Positions)
@@ -120,15 +176,16 @@ Save(ByRef Positions)
   
   Loop, %id%
   {
-    DoIgnore := DoIgnoreSize := false
-    p := {}
     i := id%A_Index%
+    p := {}
+    DoIgnore := DoIgnoreSize := false
+    
     WinReallyGetPos(i, x, y, w, h)
+    WinGet, exe, ProcessName, ahk_id %i%
     SysGet, bx, 32 ; window border thickness (x)
     SysGet, by, 33 ; window border thickness (y)
     cw := w - bx * 2 ; client width
     ch := h - by * 2 ; client height
-    WinGet, exe, ProcessName, ahk_id %i%
     
     ; Custom rules