From a63a3e710fc37b235c468e93f9a12c6349cb6c5f Mon Sep 17 00:00:00 2001
From: Albert Zhang <alchzh@gmail.com>
Date: Thu, 30 Nov 2023 00:39:06 -0500
Subject: [PATCH] Add native param in setBadgeCount() and fix settings menu

---
 src/main/appBadge.ts                 |  4 +-
 src/main/ipc.ts                      |  4 +-
 src/preload/VesktopNative.ts         |  3 +-
 src/renderer/appBadge.ts             |  2 +-
 src/renderer/components/Settings.tsx | 64 ++++++++++++++++++----------
 5 files changed, 50 insertions(+), 27 deletions(-)

diff --git a/src/main/appBadge.ts b/src/main/appBadge.ts
index c462c1e..9ba6b67 100644
--- a/src/main/appBadge.ts
+++ b/src/main/appBadge.ts
@@ -31,13 +31,15 @@ function loadTrayIcon(index: number) {
 
 let lastIndex: null | number = -1;
 
-export function setBadgeCount(count: number, tray: boolean = false) {
+export function setBadgeCount(count: number, native: boolean = true, tray: boolean = false) {
     const [index, description] = getBadgeIndexAndDescription(count);
 
     if (tray) {
         globals.tray?.setImage(loadTrayIcon(index ?? 0));
     }
 
+    if (!native) return;
+
     switch (process.platform) {
         case "linux":
             if (count === -1) count = 0;
diff --git a/src/main/ipc.ts b/src/main/ipc.ts
index 9e77c85..192fe25 100644
--- a/src/main/ipc.ts
+++ b/src/main/ipc.ts
@@ -117,7 +117,9 @@ handle(IpcEvents.SELECT_VENCORD_DIR, async () => {
     return dir;
 });
 
-handle(IpcEvents.SET_BADGE_COUNT, (_, count: number, tray: boolean) => setBadgeCount(count, tray));
+handle(IpcEvents.SET_BADGE_COUNT, (_, count: number, native: boolean, tray: boolean) =>
+    setBadgeCount(count, native, tray)
+);
 
 function readCss() {
     return readFile(VENCORD_QUICKCSS_FILE, "utf-8").catch(() => "");
diff --git a/src/preload/VesktopNative.ts b/src/preload/VesktopNative.ts
index c5569b8..431d52e 100644
--- a/src/preload/VesktopNative.ts
+++ b/src/preload/VesktopNative.ts
@@ -23,7 +23,8 @@ export const VesktopNative = {
     app: {
         relaunch: () => invoke<void>(IpcEvents.RELAUNCH),
         getVersion: () => sendSync<void>(IpcEvents.GET_VERSION),
-        setBadgeCount: (count: number, tray: boolean = false) => invoke<void>(IpcEvents.SET_BADGE_COUNT, count, tray),
+        setBadgeCount: (count: number, native: boolean = true, tray: boolean = false) =>
+            invoke<void>(IpcEvents.SET_BADGE_COUNT, count, native, tray),
         supportsWindowsTransparency: () => sendSync<boolean>(IpcEvents.SUPPORTS_WINDOWS_TRANSPARENCY)
     },
     autostart: {
diff --git a/src/renderer/appBadge.ts b/src/renderer/appBadge.ts
index be2cbb3..8eec4cc 100644
--- a/src/renderer/appBadge.ts
+++ b/src/renderer/appBadge.ts
@@ -23,7 +23,7 @@ export function setBadge() {
         if (!totalCount && hasUnread && !disableUnreadBadge) totalCount = -1;
 
         if (Settings.store.appBadge || Settings.store.trayBadge)
-            VesktopNative.app.setBadgeCount(totalCount, Settings.store.trayBadge);
+            VesktopNative.app.setBadgeCount(totalCount, Settings.store.appBadge, Settings.store.trayBadge);
     } catch (e) {
         console.error(e);
     }
diff --git a/src/renderer/components/Settings.tsx b/src/renderer/components/Settings.tsx
index b020768..b3cc18f 100644
--- a/src/renderer/components/Settings.tsx
+++ b/src/renderer/components/Settings.tsx
@@ -28,26 +28,6 @@ export default function SettingsUi() {
             "Discord Titlebar",
             "Use Discord's custom title bar instead of the Windows one. Requires a full restart."
         ],
-        !isMac && ["tray", "Tray Icon", "Add a tray icon for Vesktop", true],
-        !isMac && [
-            "minimizeToTray",
-            "Minimize to tray",
-            "Hitting X will make Vesktop minimize to the tray instead of closing",
-            true,
-            () => Settings.tray ?? true
-        ],
-        !isMac && [
-            "trayBadge",
-            "Tray Notification Badge",
-            "Show mention badge on the tray icon",
-            false,
-            () => Settings.tray ?? true,
-            v => {
-                Settings.trayBadge = v;
-                if (v) setBadge();
-                else VesktopNative.app.setBadgeCount(0, true);
-            }
-        ],
         ["arRPC", "Rich Presence", "Enables Rich Presence via arRPC", false],
         [
             "disableMinSize",
@@ -107,18 +87,56 @@ export default function SettingsUi() {
                 onChange={v => {
                     Settings.appBadge = v;
                     if (v) setBadge();
-                    else VesktopNative.app.setBadgeCount(0, Settings.trayBadge);
+                    else VesktopNative.app.setBadgeCount(0, true, false);
                 }}
                 note="Show mention badge on the app (taskbar/panel) icon"
             >
                 Notification Badge
             </Switch>
 
-            {switches.map(([key, text, note, def, predicate, onChange]) => (
+            {!isMac && (
+                <>
+                    <Switch
+                        value={Settings.tray ?? true}
+                        onChange={v => {
+                            Settings.tray = v;
+                            if (v && Settings.trayBadge) setBadge();
+                        }}
+                        note="Add a tray icon for Vesktop"
+                        key="tray"
+                    >
+                        Tray Icon
+                    </Switch>
+                    <Switch
+                        value={Settings.minimizeToTray ?? true}
+                        onChange={v => (Settings.minimizeToTray = v)}
+                        disabled={!(Settings.tray ?? true)}
+                        note="Hitting X will make Vesktop minimize to the tray instead of closing"
+                        key="minimizeToTray"
+                    >
+                        Minimize to Tray
+                    </Switch>
+                    <Switch
+                        value={Settings.trayBadge ?? true}
+                        onChange={v => {
+                            Settings.trayBadge = v;
+                            if (v) setBadge();
+                            else VesktopNative.app.setBadgeCount(0, false, true);
+                        }}
+                        disabled={!(Settings.tray ?? true)}
+                        note="Show mention badge on the tray icon"
+                        key="trayBadge"
+                    >
+                        Tray Notification Badge
+                    </Switch>
+                </>
+            )}
+
+            {switches.map(([key, text, note, def, predicate]) => (
                 <Switch
                     value={(Settings[key as any] ?? def ?? false) && predicate?.() !== false}
                     disabled={predicate && !predicate()}
-                    onChange={onChange ?? (v => (Settings[key as any] = v))}
+                    onChange={v => (Settings[key as any] = v)}
                     note={note}
                     key={key}
                 >