From 02ca35130d3244bf10298f9b2fd2af83ad4d41a3 Mon Sep 17 00:00:00 2001
From: zhiayang <gpg@tsumaran.ai>
Date: Fri, 23 Feb 2024 01:27:32 -0500
Subject: [PATCH] (re)implement charging state support

---
 app/include/zmk/battery.h                     |  1 +
 .../zmk/events/battery_state_changed.h        |  4 ++-
 app/src/activity.c                            |  3 +-
 app/src/battery.c                             | 33 +++++++++++++++++--
 4 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/app/include/zmk/battery.h b/app/include/zmk/battery.h
index edc8fd7a..8d7abc19 100644
--- a/app/include/zmk/battery.h
+++ b/app/include/zmk/battery.h
@@ -7,3 +7,4 @@
 #pragma once
 
 uint8_t zmk_battery_state_of_charge(void);
+bool zmk_battery_is_charging(void);
diff --git a/app/include/zmk/events/battery_state_changed.h b/app/include/zmk/events/battery_state_changed.h
index 157490d9..fe0d3340 100644
--- a/app/include/zmk/events/battery_state_changed.h
+++ b/app/include/zmk/events/battery_state_changed.h
@@ -12,6 +12,7 @@
 struct zmk_battery_state_changed {
     // TODO: Other battery channels
     uint8_t state_of_charge;
+    bool is_charging;
 };
 
 ZMK_EVENT_DECLARE(zmk_battery_state_changed);
@@ -20,6 +21,7 @@ struct zmk_peripheral_battery_state_changed {
     uint8_t source;
     // TODO: Other battery channels
     uint8_t state_of_charge;
+    bool is_charging;
 };
 
-ZMK_EVENT_DECLARE(zmk_peripheral_battery_state_changed);
\ No newline at end of file
+ZMK_EVENT_DECLARE(zmk_peripheral_battery_state_changed);
diff --git a/app/src/activity.c b/app/src/activity.c
index 8f421f85..ba9ee5f9 100644
--- a/app/src/activity.c
+++ b/app/src/activity.c
@@ -20,6 +20,7 @@ LOG_MODULE_DECLARE(zmk, CONFIG_ZMK_LOG_LEVEL);
 #include <zmk/events/position_state_changed.h>
 #include <zmk/events/sensor_event.h>
 
+#include <zmk/battery.h>
 #include <zmk/activity.h>
 
 #if IS_ENABLED(CONFIG_USB_DEVICE_STACK)
@@ -126,7 +127,7 @@ void activity_work_handler(struct k_work *work) {
     int32_t current = k_uptime_get();
     int32_t inactive_time = current - activity_last_uptime;
 #if IS_ENABLED(CONFIG_ZMK_SLEEP)
-    if (inactive_time > MAX_SLEEP_MS && !is_usb_power_present()) {
+    if (inactive_time > MAX_SLEEP_MS && !(is_usb_power_present() || zmk_battery_is_charging())) {
         // Put devices in suspend power mode before sleeping
         set_state(ZMK_ACTIVITY_SLEEP);
 
diff --git a/app/src/battery.c b/app/src/battery.c
index 1295f822..6efb275f 100644
--- a/app/src/battery.c
+++ b/app/src/battery.c
@@ -19,12 +19,16 @@ LOG_MODULE_DECLARE(zmk, CONFIG_ZMK_LOG_LEVEL);
 #include <zmk/battery.h>
 #include <zmk/events/battery_state_changed.h>
 #include <zmk/events/activity_state_changed.h>
+#include <zmk/events/usb_conn_state_changed.h>
 #include <zmk/activity.h>
 #include <zmk/workqueue.h>
+#include <zmk/usb.h>
 
 static uint8_t last_state_of_charge = 0;
+static bool last_battery_is_charging = false;
 
 uint8_t zmk_battery_state_of_charge(void) { return last_state_of_charge; }
+bool zmk_battery_is_charging(void) { return last_battery_is_charging; }
 
 #if DT_HAS_CHOSEN(zmk_battery)
 static const struct device *const battery = DEVICE_DT_GET(DT_CHOSEN(zmk_battery));
@@ -51,8 +55,19 @@ static int zmk_battery_update(const struct device *battery) {
         return rc;
     }
 
-    if (last_state_of_charge != state_of_charge.val1) {
+    // TODO: for now, battery charging is determined solely by USB being plugged in
+#if IS_ENABLED(CONFIG_USB_DEVICE_STACK)
+    const bool batt_is_charging = zmk_usb_is_powered();
+#else
+    const bool batt_is_charging = false;
+#endif
+
+    if (last_state_of_charge != state_of_charge.val1 ||
+        last_battery_is_charging != batt_is_charging) {
+
         last_state_of_charge = state_of_charge.val1;
+        last_battery_is_charging = batt_is_charging;
+
 #if IS_ENABLED(CONFIG_BT_BAS)
         LOG_DBG("Setting BAS GATT battery level to %d.", last_state_of_charge);
 
@@ -63,8 +78,10 @@ static int zmk_battery_update(const struct device *battery) {
             return rc;
         }
 #endif
-        rc = raise_zmk_battery_state_changed(
-            (struct zmk_battery_state_changed){.state_of_charge = last_state_of_charge});
+        rc = raise_zmk_battery_state_changed((struct zmk_battery_state_changed){
+            .state_of_charge = last_state_of_charge,
+            .is_charging = batt_is_charging,
+        });
     }
 
     return rc;
@@ -127,6 +144,12 @@ static int battery_event_listener(const zmk_event_t *eh) {
             break;
         }
     }
+#if IS_ENABLED(CONFIG_USB_DEVICE_STACK)
+    else if (as_zmk_usb_conn_state_changed(eh)) {
+        // update the battery on the workqueue if usb connection changed.
+        k_work_submit_to_queue(zmk_workqueue_lowprio_work_q(), &battery_work);
+    }
+#endif
     return -ENOTSUP;
 }
 
@@ -134,4 +157,8 @@ ZMK_LISTENER(battery, battery_event_listener);
 
 ZMK_SUBSCRIPTION(battery, zmk_activity_state_changed);
 
+#if IS_ENABLED(CONFIG_USB_DEVICE_STACK)
+ZMK_SUBSCRIPTION(battery, zmk_usb_conn_state_changed);
+#endif
+
 SYS_INIT(zmk_battery_init, APPLICATION, CONFIG_APPLICATION_INIT_PRIORITY);