From 1e035f1a3646c9caa76c1c8f809f4e0269cb228b Mon Sep 17 00:00:00 2001
From: DevilXD <bartosz.peszek@gmail.com>
Date: Sat, 3 Sep 2022 18:12:32 +0200
Subject: [PATCH] Properly handle Windows shutdown message

---
 gui.py           | 38 ++++++++++++++++++++++++++++++++++----
 main.py          |  7 +++++--
 requirements.txt |  1 +
 3 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/gui.py b/gui.py
index 5520536..ca383e4 100644
--- a/gui.py
+++ b/gui.py
@@ -5,14 +5,17 @@ import logging
 import webbrowser
 import tkinter as tk
 from math import log10, ceil
-from functools import partial
 from collections import abc, namedtuple
 from tkinter.font import Font, nametofont
+from functools import partial, cached_property
 from datetime import datetime, timedelta, timezone
 from typing import Any, TypedDict, NoReturn, Generic, TYPE_CHECKING
 from tkinter import Tk, ttk, StringVar, DoubleVar, IntVar, PhotoImage
 
 import pystray
+import win32api
+import win32con
+import win32gui
 from yarl import URL
 from PIL import Image as Image_module
 
@@ -1631,8 +1634,6 @@ class GUIManager:
         # root.resizable(False, True)
         root.iconbitmap(resource_path("pickaxe.ico"))  # window icon
         root.title(WINDOW_TITLE)  # window title
-        root.protocol("WM_DELETE_WINDOW", self.close)  # hook the X window closing button
-        root.protocol("WM_ENDSESSION", self.close)  # hook the Windows shutdown signal
         root.bind_all("<KeyPress-Escape>", self.unfocus)  # pressing ESC unfocuses selection
         # Image cache for displaying images
         self._cache = ImageCache(self)
@@ -1726,6 +1727,18 @@ class GUIManager:
             self._root.after_idle(self.tray.minimize)
         else:
             self._root.deiconify()
+        # NOTE: this root.update() is required for the below to work - don't remove
+        root.update()
+        self._message_map = {
+            # window close request
+            win32con.WM_CLOSE: self.close,
+            # shutdown request
+            win32con.WM_QUERYENDSESSION: self.close,
+        }
+        # This hooks up the wnd_proc function as the message processor for the root window.
+        self.old_wnd_proc = win32gui.SetWindowLong(
+            self._handle, win32con.GWL_WNDPROC, self.wnd_proc
+        )
 
     # https://stackoverflow.com/questions/56329342/tkinter-treeview-background-tag-not-working
     def _fixed_map(self, option):
@@ -1742,6 +1755,21 @@ class GUIManager:
             if elm[:2] != ("!disabled", "!selected")
         ]
 
+    def wnd_proc(self, hwnd, msg, w_param, l_param):
+        """
+        This function serves as a message processor for all messages sent
+        to the application by Windows.
+        """
+        if msg == win32con.WM_DESTROY:
+            win32api.SetWindowLong(self._handle, win32con.GWL_WNDPROC, self.old_wnd_proc)
+        if msg in self._message_map:
+            return self._message_map[msg](w_param, l_param)
+        return win32gui.CallWindowProc(self.old_wnd_proc, hwnd, msg, w_param, l_param)
+
+    @cached_property
+    def _handle(self) -> int:
+        return int(self._root.wm_frame(), 16)
+
     @property
     def running(self) -> bool:
         return self._poll_task is not None
@@ -1785,7 +1813,7 @@ class GUIManager:
             await asyncio.sleep(0.05)
         self._poll_task = None
 
-    def close(self, *args):
+    def close(self, *args) -> int:
         """
         Requests the GUI application to close.
         The window itself will be closed in the closing sequence later.
@@ -1793,6 +1821,7 @@ class GUIManager:
         self._close_requested.set()
         # notify client we're supposed to close
         self._twitch.close()
+        return 0
 
     def close_window(self):
         """
@@ -1828,6 +1857,7 @@ class GUIManager:
 
 if __name__ == "__main__":
     # Everything below is for debug purposes only
+    import aiohttp
     from types import SimpleNamespace
 
     class StrNamespace(SimpleNamespace):
diff --git a/main.py b/main.py
index f62ecc8..858d46b 100644
--- a/main.py
+++ b/main.py
@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import io
 import sys
-import ctypes
 import signal
 import asyncio
 import logging
@@ -25,6 +24,10 @@ try:
     import PIL  # noqa
 except ModuleNotFoundError as exc:
     raise ImportError("You have to run 'pip install pillow' first") from exc
+try:
+    import win32gui  # noqa
+except ModuleNotFoundError as exc:
+    raise ImportError("You have to run 'pip install pywin32' first") from exc
 
 from translate import _
 from twitch import Twitch
@@ -126,7 +129,7 @@ root.destroy()
 
 # check if we're not already running
 try:
-    exists = ctypes.windll.user32.FindWindowW(None, WINDOW_TITLE)
+    exists = win32gui.FindWindow(None, WINDOW_TITLE)
 except AttributeError:
     # we're not on Windows - continue
     exists = False
diff --git a/requirements.txt b/requirements.txt
index b3cbb08..4eebaed 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
 aiohttp>2.0,<4.0
 Pillow
 pystray
+pywin32