From 490b8c8b2f415461124658cc60c7ccd1316bdee1 Mon Sep 17 00:00:00 2001
From: Oleh Polisan <polisanoleg@gmail.com>
Date: Thu, 20 Jun 2024 16:01:03 +0300
Subject: [PATCH] feat: ability to set custom icons

---
 src/main/mainWindow.ts                            | 12 +++++++++---
 src/main/tray.ts                                  |  6 +++++-
 src/preload/VesktopNative.ts                      |  3 ++-
 src/renderer/components/settings/TraySettings.tsx |  9 ++++++++-
 src/renderer/patches/tray.ts                      |  2 ++
 5 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/src/main/mainWindow.ts b/src/main/mainWindow.ts
index a099e5a..374145f 100755
--- a/src/main/mainWindow.ts
+++ b/src/main/mainWindow.ts
@@ -23,7 +23,7 @@ import { isTruthy } from "shared/utils/guards";
 import { once } from "shared/utils/once";
 import type { SettingsStore } from "shared/utils/SettingsStore";
 
-import { ICON_PATH, STATIC_DIR } from "../shared/paths";
+import { ICON_PATH } from "../shared/paths";
 import { createAboutWindow } from "./about";
 import { initArRPC } from "./arrpc";
 import {
@@ -125,6 +125,12 @@ function initTray(win: BrowserWindow) {
     ]);
 
     tray = new Tray(ICON_PATH);
+    try {
+        if (Settings.store.trayCustom) tray.setImage(join(DATA_DIR, "TrayIcons", "icon.png"));
+    } catch (error) {
+        console.log("Error while loading custom tray image. Recreating new ones.");
+        generateTrayIcons(true);
+    }
     tray.setToolTip("Vesktop");
     tray.setContextMenu(trayMenu);
     tray.on("click", onTrayClick);
@@ -505,7 +511,8 @@ export async function createWindows() {
 
 export async function setTrayIcon(iconName: string) {
     if (!tray || tray.isDestroyed()) return;
-    if (iconName !== "icon") {
+    const Icons = new Set(["speaking", "muted", "deafened", "idle", "icon"]);
+    if (Icons.has(iconName)) {
         try {
             tray.setImage(join(DATA_DIR, "TrayIcons", iconName + ".png"));
         } catch (error) {
@@ -514,5 +521,4 @@ export async function setTrayIcon(iconName: string) {
         }
         return;
     }
-    tray.setImage(join(STATIC_DIR, "icon.png"));
 }
diff --git a/src/main/tray.ts b/src/main/tray.ts
index c2f8ecc..a8facd5 100644
--- a/src/main/tray.ts
+++ b/src/main/tray.ts
@@ -53,11 +53,15 @@ export async function generateTrayIcons(force = false) {
         for (const icon of Icons) {
             mainWin.webContents.send(IpcEvents.CREATE_TRAY_ICON_REQUEST, icon);
         }
+        copyFileSync(join(STATIC_DIR, "icon.png"), join(DATA_DIR, "TrayIcons", "icon.png"));
+        mainWin.webContents.send(IpcEvents.SET_CURRENT_VOICE_TRAY_ICON);
     }
-    mainWin.webContents.send(IpcEvents.SET_CURRENT_VOICE_TRAY_ICON);
 }
 
 export async function pickTrayIcon(iconName: string) {
+    const Icons = new Set(["speaking", "muted", "deafened", "idle", "icon"]);
+    if (!Icons.has(iconName)) return;
+
     const res = await dialog.showOpenDialog(mainWin!, {
         properties: ["openFile"],
         filters: [{ name: "Image", extensions: ["png", "jpg"] }]
diff --git a/src/preload/VesktopNative.ts b/src/preload/VesktopNative.ts
index 53d6372..a60005f 100644
--- a/src/preload/VesktopNative.ts
+++ b/src/preload/VesktopNative.ts
@@ -36,7 +36,8 @@ export const VesktopNative = {
     fileManager: {
         showItemInFolder: (path: string) => invoke<void>(IpcEvents.SHOW_ITEM_IN_FOLDER, path),
         selectVencordDir: () => invoke<LiteralUnion<"cancelled" | "invalid", string>>(IpcEvents.SELECT_VENCORD_DIR),
-        selectTrayIcon: () => invoke<LiteralUnion<"cancelled" | "invalid", string>>(IpcEvents.SELECT_TRAY_ICON)
+        selectTrayIcon: (iconName: string) =>
+            invoke<LiteralUnion<"cancelled" | "invalid", string>>(IpcEvents.SELECT_TRAY_ICON, iconName)
     },
     settings: {
         get: () => sendSync<Settings>(IpcEvents.GET_SETTINGS),
diff --git a/src/renderer/components/settings/TraySettings.tsx b/src/renderer/components/settings/TraySettings.tsx
index bc5a88b..8439094 100644
--- a/src/renderer/components/settings/TraySettings.tsx
+++ b/src/renderer/components/settings/TraySettings.tsx
@@ -44,13 +44,14 @@ function trayEditButton(iconName: string) {
                 alt="read if cute :3"
                 width="48"
                 height="48"
+                id={iconName}
             ></img>
             <PencilIcon
                 className="vcd-edit-button"
                 width="40"
                 height="40"
                 onClick={async () => {
-                    const choice = await VesktopNative.fileManager.selectTrayIcon();
+                    const choice = await VesktopNative.fileManager.selectTrayIcon(iconName);
                     switch (choice) {
                         case "cancelled":
                             return;
@@ -65,6 +66,12 @@ function trayEditButton(iconName: string) {
                     console.log("choice:", choice);
                     // copy image and reload
                     // settings.trayIconPath = choice;
+                    const iconDataURL = VesktopNative.tray.getIconSync(iconName);
+                    const img = document.getElementById(iconName) as HTMLImageElement;
+                    if (img) {
+                        img.src = iconDataURL;
+                    }
+                    VesktopNative.tray.createIconResponse(iconName, iconDataURL);
                 }}
             />
         </div>
diff --git a/src/renderer/patches/tray.ts b/src/renderer/patches/tray.ts
index 703676d..be3f46f 100644
--- a/src/renderer/patches/tray.ts
+++ b/src/renderer/patches/tray.ts
@@ -22,6 +22,8 @@ export function setCurrentTrayIcon() {
         } else {
             VesktopNative.tray.setIcon("idle");
         }
+    } else {
+        VesktopNative.tray.setIcon("icon");
     }
 }