From b1c9691b93b5235736a2523d6b67f1dca6a29f61 Mon Sep 17 00:00:00 2001
From: Chi Huu Huynh <73843190+Chi-EEE@users.noreply.github.com>
Date: Fri, 15 Mar 2024 08:21:47 +0000
Subject: [PATCH] Add more methods for Electron

---
 app/admin_panel/main.js    | 51 +++++++++++++++++++++++++++++++++-----
 app/admin_panel/preload.js |  5 +++-
 2 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/app/admin_panel/main.js b/app/admin_panel/main.js
index b004bb64..519a8eba 100644
--- a/app/admin_panel/main.js
+++ b/app/admin_panel/main.js
@@ -115,6 +115,9 @@ class WebSocketServer {
         /** @type {WebSocket.Server | undefined} */
         this._wss = undefined;
 
+        /** @type {number | undefined} */
+        this._port = undefined;
+
         /** @type {Code} */
         this._code = new Code();
 
@@ -131,6 +134,7 @@ class WebSocketServer {
      */
     connect(port) {
         this._wss = new WebSocket.Server({ port: port });
+        this._port = port;
     }
 
     /**
@@ -142,11 +146,15 @@ class WebSocketServer {
         }
     }
 
-    code() {
+    generateCode() {
         this._code.generate();
         return this._code.get();
     }
 
+    getCode() {
+        return this._code.get();
+    }
+
     /**
      * This function only allows a single connection to the WebSocket server.
      */
@@ -156,7 +164,7 @@ class WebSocketServer {
             // console.log(req.socket.remoteAddress)
             const uuid = crypto.randomUUID();
             this._raspberry_pi_clients.set(uuid, ws);
-           
+
             ws.send(JSON.stringify({ uuid: uuid }));
 
             mainWindow.webContents.send('onConnection', JSON.stringify({ uuid: uuid }));
@@ -169,7 +177,7 @@ class WebSocketServer {
             });
 
             ws.on('message', async (message) => {
-                if (this._client === ws) {
+                if (uuid !== this._client.uuid) {
                     return;
                 }
                 mainWindow.webContents.send('onMessage', message.toString());
@@ -189,7 +197,7 @@ async function startWebSocketServer(_event, args) {
     try {
         websocket_server.connect(args.port);
         websocket_server.createConnection();
-        return { success: true, code: websocket_server.code() };
+        return { success: true, code: websocket_server.generateCode() };
     } catch (error) {
         return { success: false, message: `Unable to start WebSocket Server, Error: ${error}` };
     }
@@ -199,11 +207,34 @@ function closeWebSocketServer(_event, _args) {
     websocket_server.close();
 }
 
+function getWebSocketServer(_event, _args) {
+    if (websocket_server._wss !== undefined) {
+        return { success: true, code: websocket_server.getCode(), port: websocket_server._port };
+    }
+    else {
+        return { success: false };
+    }
+}
+
 function onClose() {
     closeWebSocketServer();
 }
 
-function connectToRaspberryPi(_event, args) {
+function connectRaspberryPi(_event, args) {
+    const uuid = args.uuid;
+    if (websocket_server._raspberry_pi_clients.has(uuid)) {
+        websocket_server._client = { uuid: uuid, ws: websocket_server._raspberry_pi_clients.get(uuid) };
+        return { success: true, message: `Connected to Raspberry Pi with UUID: ${uuid}` };
+    } else {
+        return { success: false, message: `Raspberry Pi with UUID: ${uuid} not found` };
+    }
+}
+
+function disconnectRaspberryPi(_event, _args) {
+    websocket_server._client = undefined;
+}
+
+function connectRaspberryPi(_event, args) {
     const uuid = args.uuid;
     if (websocket_server._raspberry_pi_clients.has(uuid)) {
         websocket_server._client = websocket_server._raspberry_pi_clients.get(uuid);
@@ -213,9 +244,17 @@ function connectToRaspberryPi(_event, args) {
     }
 }
 
+function getRaspberryPi(_event, _args) {
+    return { uuid: websocket_server._client.uuid }
+}
 
 app.on('window-all-closed', onClose);
 ipcMain.handle('getLocalIPList', getLocalIPList);
+
 ipcMain.handle('startWebSocketServer', startWebSocketServer);
 ipcMain.handle('closeWebSocketServer', closeWebSocketServer);
-ipcMain.handle('connectToRaspberryPi', connectToRaspberryPi);
+ipcMain.handle('getWebSocketServer', getWebSocketServer);
+
+ipcMain.handle('connectRaspberryPi', connectRaspberryPi);
+ipcMain.handle('disconnectRaspberryPi', disconnectRaspberryPi);
+ipcMain.handle('getRaspberryPi', getRaspberryPi);
diff --git a/app/admin_panel/preload.js b/app/admin_panel/preload.js
index 55e16774..7dba4e09 100644
--- a/app/admin_panel/preload.js
+++ b/app/admin_panel/preload.js
@@ -12,7 +12,10 @@ const WINDOW_API = {
     onConnection: (/** @type {(string) => void} */ callback) => ipcRenderer.on('onConnection', (_event, /** @type {string} */ value) => callback(value)),
     onDisconnection: (/** @type {(string) => void} */ callback) => ipcRenderer.on('onDisconnection', (_event, /** @type {string} */ value) => callback(value)),
 
-    connectToRaspberryPi: (/** @type {string} */ uuid) => ipcRenderer.invoke("connectToRaspberryPi", uuid),
+    connectRaspberryPi: (/** @type {string} */ uuid) => ipcRenderer.invoke("connectRaspberryPi", uuid),
+    disconnectRaspberryPi: () => ipcRenderer.invoke("disconnectRaspberryPi"),
+    
+    getRaspberryPi: () => ipcRenderer.invoke("disconnectRaspberryPi"),
 }
 
 contextBridge.exposeInMainWorld("api", WINDOW_API)
\ No newline at end of file