From 06015646e9ad832a5c3afcde1b61984c0f45a5d9 Mon Sep 17 00:00:00 2001
From: lennard <lennardkittner@icloud.com>
Date: Sat, 10 Jun 2023 12:18:06 +0200
Subject: [PATCH] More error handling.

---
 src/battery_tray.rs | 26 +++++++++++++-------
 src/lib.rs          | 10 ++++----
 src/main.rs         | 58 ++++++++++++++++++++++++---------------------
 3 files changed, 54 insertions(+), 40 deletions(-)

diff --git a/src/battery_tray.rs b/src/battery_tray.rs
index 59dd802..a8afac3 100644
--- a/src/battery_tray.rs
+++ b/src/battery_tray.rs
@@ -4,20 +4,27 @@ use ksni::{Tray, MenuItem, menu::{StandardItem}, ToolTip};
 pub struct BatteryTray {
     battery_level: u8,
     charging: bool,
+    device_found: bool,
 }
 
 impl BatteryTray {
-    pub fn new(battery_level: u8, charging: bool) -> Self {
+    pub fn new() -> Self {
         BatteryTray {
-            battery_level,
-            charging,
+            battery_level: 0,
+            charging: false,
+            device_found: false,
         }
     }
 
     pub fn update(&mut self, battery_level: u8, charging: bool) {
+        self.device_found = true;
         self.battery_level = battery_level;
         self.charging = charging;
     }
+
+    pub fn no_device_found(&mut self) {
+        self.device_found = false;
+    }
 }
 
 impl Tray for BatteryTray {
@@ -36,11 +43,14 @@ impl Tray for BatteryTray {
         ]
     }
     fn tool_tip(&self) -> ToolTip {
-        let description = if self.charging {
-            format!("Battery level: {}%\nCharging", self.battery_level)
-        } else {
-            format!("Battery level: {}%\nNot charging", self.battery_level)
-        };
+        let description = 
+            if !self.device_found {
+                "No device found".to_string()
+            } else if self.charging {
+                format!("Battery level: {}%\nCharging", self.battery_level)
+            } else {
+                format!("Battery level: {}%\nNot charging", self.battery_level)
+            };
         ToolTip {
             title: "HyperX Cloud II".to_string(),
             description: description,
diff --git a/src/lib.rs b/src/lib.rs
index c6f9d4f..cc52295 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -26,8 +26,8 @@ pub enum DeviceError {
     NoDeviceFound(),
     #[error("Error: No response. Is the headset turned on?")]
     HeadSetOff(),
-    #[error("Error: Unknown response.")]
-    UnknownResponse(),
+    #[error("Error: Unknown response: {0:?}")]
+    UnknownResponse([u8; 8]),
 }
    
 pub struct Device {
@@ -50,17 +50,17 @@ impl Device {
     pub fn get_battery_level(&self) -> Result<(u8, bool), DeviceError> {
         self.hid_device.write(&BATTERY_PACKET)?;
         let mut buf = [0u8; 8];
-        let res = self.hid_device.read_timeout(&mut buf[..], 1000)?;
+        let res = self.hid_device.read_timeout(&mut buf[..], 800)?;
         if res == 0 {
             return Err(DeviceError::HeadSetOff());
         }
         if !buf.starts_with(&PREAMBLE) {
-            return Err(DeviceError::UnknownResponse());
+            return Err(DeviceError::UnknownResponse(buf));
         }
         let charging = match buf[CHARGING_INDEX] {
             CHARGING_STATE => true,
             NOT_CHARGING_STATE => false,
-            _ => return Err(DeviceError::UnknownResponse()),
+            _ => return Err(DeviceError::UnknownResponse(buf)),
         };
         Ok((buf[BATTERY_LEVEL_INDEX], charging))
     }
diff --git a/src/main.rs b/src/main.rs
index 3d7eea9..a84cfdb 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,46 +1,50 @@
-use hyper_x_cloud_ii_wireless::Device;
+use hyper_x_cloud_ii_wireless::{Device, DeviceError};
 use ksni::TrayService;
-use battery_tray::BatteryTray;
 mod battery_tray;
+use battery_tray::BatteryTray;
+
+fn pair_device() -> Device {
+    loop {
+        match Device::new() {
+            Ok(device) => break device,
+            Err(error) => {
+                eprintln!("{error}");
+            }
+        };
+        std::thread::sleep(std::time::Duration::from_secs(1));
+    }
+}
 
 fn main() {
-    let device = match Device::new() {
-        Ok(device) => device,
-        Err(error) => {
-            eprintln!("{error}");
-            std::process::exit(1);
-        }
-    };
-    let (battery_level, charging) = match device.get_battery_level() {
-        Ok(t) => t,
-        Err(error) => {
-            eprintln!("{error}");
-            std::process::exit(1);
-        }
-    };
-    println!("Battery level: {}%", battery_level);
-    if charging {
-        println!("Charging");
-    } else {
-        println!("Not charging");
-    }
-
-    let service = TrayService::new(BatteryTray::new(battery_level, charging));
+    let service = TrayService::new(BatteryTray::new());
     let handle = service.handle();
     service.spawn();
 
+    let mut device = pair_device();
+
     // Run loop
     loop {
         std::thread::sleep(std::time::Duration::from_secs(1));
         let (battery_level, charging) = match device.get_battery_level() {
             Ok(t) => t,
+            Err(DeviceError::HidError(hidapi::HidError::HidApiError { message })) => {
+                eprintln!("Error: {message}");
+                if message == "No such device" {
+                    handle.update(|tray: &mut BatteryTray| { tray.no_device_found(); });
+                    device = pair_device();
+                }
+                continue;
+            }
+            Err(DeviceError::NoDeviceFound()) => {
+                eprintln!("Error: {}", DeviceError::NoDeviceFound());
+                handle.update(|tray: &mut BatteryTray| { tray.no_device_found(); });
+                continue;
+            }
             Err(error) => {
                 eprintln!("{error}");
                 continue;
             }
         };
-        handle.update(|tray: &mut BatteryTray| {
-            tray.update(battery_level, charging);
-        });
+        handle.update(|tray: &mut BatteryTray| { tray.update(battery_level, charging); });
     }
 }