From 8973076465f2b97ba4c68c1eaa88b4a60ac4d3d4 Mon Sep 17 00:00:00 2001
From: Peter Johanson <peter@peterjohanson.com>
Date: Mon, 2 Dec 2024 00:27:48 -0700
Subject: [PATCH] feat(behaviors): Add retention boot mode to reset.

Support new generic Zephyr retention boot mode API in the reset
behavior.
---
 app/dts/behaviors/reset.dtsi                  |  1 +
 .../behaviors/zmk,behavior-reset.yaml         |  2 ++
 app/src/behaviors/behavior_reset.c            | 28 +++++++++++++++++--
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/app/dts/behaviors/reset.dtsi b/app/dts/behaviors/reset.dtsi
index 2aa41d7d6e8..1b29f9d19aa 100644
--- a/app/dts/behaviors/reset.dtsi
+++ b/app/dts/behaviors/reset.dtsi
@@ -19,6 +19,7 @@
         bootloader: bootload {
             compatible = "zmk,behavior-reset";
             type = <RST_UF2>;
+            bootloader;
             #binding-cells = <0>;
             display-name = "Bootloader";
         };
diff --git a/app/dts/bindings/behaviors/zmk,behavior-reset.yaml b/app/dts/bindings/behaviors/zmk,behavior-reset.yaml
index 9eb7aa6a3c4..aeb9450eea9 100644
--- a/app/dts/bindings/behaviors/zmk,behavior-reset.yaml
+++ b/app/dts/bindings/behaviors/zmk,behavior-reset.yaml
@@ -11,3 +11,5 @@ properties:
   type:
     type: int
     default: 0
+  bootloader:
+    type: boolean
diff --git a/app/src/behaviors/behavior_reset.c b/app/src/behaviors/behavior_reset.c
index 554132f4a1e..360f2eb7691 100644
--- a/app/src/behaviors/behavior_reset.c
+++ b/app/src/behaviors/behavior_reset.c
@@ -14,11 +14,21 @@
 
 #include <zmk/behavior.h>
 
+#if IS_ENABLED(CONFIG_RETENTION_BOOT_MODE)
+
+#include <zephyr/retention/bootmode.h>
+
+#endif /* IS_ENABLED(CONFIG_RETENTION_BOOT_MODE) */
+
 LOG_MODULE_DECLARE(zmk, CONFIG_ZMK_LOG_LEVEL);
 
 #if DT_HAS_COMPAT_STATUS_OKAY(DT_DRV_COMPAT)
 struct behavior_reset_config {
+#if IS_ENABLED(CONFIG_RETENTION_BOOT_MODE)
+    enum BOOT_MODE_TYPES boot_mode;
+#else
     int type;
+#endif /* IS_ENABLED(CONFIG_RETENTION_BOOT_MODE) */
 };
 
 static int behavior_reset_init(const struct device *dev) { return 0; };
@@ -28,10 +38,20 @@ static int on_keymap_binding_pressed(struct zmk_behavior_binding *binding,
     const struct device *dev = zmk_behavior_get_binding(binding->behavior_dev);
     const struct behavior_reset_config *cfg = dev->config;
 
-    // TODO: Correct magic code for going into DFU?
+#if IS_ENABLED(CONFIG_RETENTION_BOOT_MODE)
+    int ret = bootmode_set(cfg->boot_mode);
+    if (ret < 0) {
+        LOG_ERR("Failed to set the bootloader mode (%d)", ret);
+        return ZMK_BEHAVIOR_OPAQUE;
+    }
+
+    sys_reboot(SYS_REBOOT_WARM);
+#else
     // See
     // https://github.com/adafruit/Adafruit_nRF52_Bootloader/blob/d6b28e66053eea467166f44875e3c7ec741cb471/src/main.c#L107
     sys_reboot(cfg->type);
+#endif /* IS_ENABLED(CONFIG_RETENTION_BOOT_MODE) */
+
     return ZMK_BEHAVIOR_OPAQUE;
 }
 
@@ -45,7 +65,11 @@ static const struct behavior_driver_api behavior_reset_driver_api = {
 
 #define RST_INST(n)                                                                                \
     static const struct behavior_reset_config behavior_reset_config_##n = {                        \
-        .type = DT_INST_PROP(n, type)};                                                            \
+        COND_CODE_1(                                                                               \
+            IS_ENABLED(CONFIG_RETENTION_BOOT_MODE),                                                \
+            (DT_INST_PROP(n, bootloader) ? BOOT_MODE_TYPE_BOOTLOADER : BOOT_MODE_TYPE_NORMAL),     \
+            (.type = DT_INST_PROP(n, type))),                                                      \
+    };                                                                                             \
     BEHAVIOR_DT_INST_DEFINE(n, behavior_reset_init, NULL, NULL, &behavior_reset_config_##n,        \
                             POST_KERNEL, CONFIG_KERNEL_INIT_PRIORITY_DEFAULT,                      \
                             &behavior_reset_driver_api);