From 339fb127d00f2ea897a2588c06b388639126ce84 Mon Sep 17 00:00:00 2001
From: Peter Johanson <peter@peterjohanson.com>
Date: Wed, 31 Jul 2024 11:54:15 -0600
Subject: [PATCH] feat: Split physical layout selection sync.

* Ensure the split peripherals have the same selected physical
  layout on connection and change.
---
 app/include/zmk/physical_layouts.h     |  7 +++
 app/include/zmk/split/bluetooth/uuid.h |  1 +
 app/src/physical_layouts.c             | 11 +++-
 app/src/split/bluetooth/central.c      | 82 +++++++++++++++++++++++++-
 app/src/split/bluetooth/service.c      | 43 +++++++++++++-
 5 files changed, 141 insertions(+), 3 deletions(-)

diff --git a/app/include/zmk/physical_layouts.h b/app/include/zmk/physical_layouts.h
index e78602e3811..33004af738f 100644
--- a/app/include/zmk/physical_layouts.h
+++ b/app/include/zmk/physical_layouts.h
@@ -8,6 +8,13 @@
 
 #include <zephyr/kernel.h>
 #include <zmk/matrix_transform.h>
+#include <zmk/event_manager.h>
+
+struct zmk_physical_layout_selection_changed {
+    uint8_t selection;
+};
+
+ZMK_EVENT_DECLARE(zmk_physical_layout_selection_changed);
 
 struct zmk_key_physical_attrs {
     int16_t width;
diff --git a/app/include/zmk/split/bluetooth/uuid.h b/app/include/zmk/split/bluetooth/uuid.h
index dccdfc804c5..4a653c73b83 100644
--- a/app/include/zmk/split/bluetooth/uuid.h
+++ b/app/include/zmk/split/bluetooth/uuid.h
@@ -18,3 +18,4 @@
 #define ZMK_SPLIT_BT_CHAR_RUN_BEHAVIOR_UUID ZMK_BT_SPLIT_UUID(0x00000002)
 #define ZMK_SPLIT_BT_CHAR_SENSOR_STATE_UUID ZMK_BT_SPLIT_UUID(0x00000003)
 #define ZMK_SPLIT_BT_UPDATE_HID_INDICATORS_UUID ZMK_BT_SPLIT_UUID(0x00000004)
+#define ZMK_SPLIT_BT_SELECT_PHYS_LAYOUT_UUID ZMK_BT_SPLIT_UUID(0x00000005)
diff --git a/app/src/physical_layouts.c b/app/src/physical_layouts.c
index ac627f46533..cc5f255df5c 100644
--- a/app/src/physical_layouts.c
+++ b/app/src/physical_layouts.c
@@ -22,6 +22,8 @@ LOG_MODULE_DECLARE(zmk, CONFIG_ZMK_LOG_LEVEL);
 #include <zmk/event_manager.h>
 #include <zmk/events/position_state_changed.h>
 
+ZMK_EVENT_IMPL(zmk_physical_layout_selection_changed);
+
 #define DT_DRV_COMPAT zmk_physical_layout
 
 #if DT_HAS_COMPAT_STATUS_OKAY(DT_DRV_COMPAT)
@@ -229,7 +231,14 @@ int zmk_physical_layouts_select(uint8_t index) {
         return -EINVAL;
     }
 
-    return zmk_physical_layouts_select_layout(layouts[index]);
+    int ret = zmk_physical_layouts_select_layout(layouts[index]);
+
+    if (ret >= 0) {
+        raise_zmk_physical_layout_selection_changed(
+            (struct zmk_physical_layout_selection_changed){.selection = index});
+    }
+
+    return ret;
 }
 
 int zmk_physical_layouts_get_selected(void) {
diff --git a/app/src/split/bluetooth/central.c b/app/src/split/bluetooth/central.c
index 0f4cd78b531..38ea243371d 100644
--- a/app/src/split/bluetooth/central.c
+++ b/app/src/split/bluetooth/central.c
@@ -30,6 +30,7 @@ LOG_MODULE_DECLARE(zmk, CONFIG_ZMK_LOG_LEVEL);
 #include <zmk/events/sensor_event.h>
 #include <zmk/events/battery_state_changed.h>
 #include <zmk/hid_indicators_types.h>
+#include <zmk/physical_layouts.h>
 
 static int start_scanning(void);
 
@@ -56,6 +57,7 @@ struct peripheral_slot {
 #if IS_ENABLED(CONFIG_ZMK_SPLIT_PERIPHERAL_HID_INDICATORS)
     uint16_t update_hid_indicators;
 #endif // IS_ENABLED(CONFIG_ZMK_SPLIT_PERIPHERAL_HID_INDICATORS)
+    uint16_t selected_physical_layout_handle;
     uint8_t position_state[POSITION_STATE_DATA_LEN];
     uint8_t changed_positions[POSITION_STATE_DATA_LEN];
 };
@@ -141,6 +143,7 @@ int release_peripheral_slot(int index) {
     // Clean up previously discovered handles;
     slot->subscribe_params.value_handle = 0;
     slot->run_behavior_handle = 0;
+    slot->selected_physical_layout_handle = 0;
 #if IS_ENABLED(CONFIG_ZMK_SPLIT_PERIPHERAL_HID_INDICATORS)
     slot->update_hid_indicators = 0;
 #endif // IS_ENABLED(CONFIG_ZMK_SPLIT_PERIPHERAL_HID_INDICATORS)
@@ -392,6 +395,46 @@ static int split_central_subscribe(struct bt_conn *conn, struct bt_gatt_subscrib
     return err;
 }
 
+static int update_peripheral_selected_layout(struct peripheral_slot *slot, uint8_t layout_idx) {
+    if (slot->state != PERIPHERAL_SLOT_STATE_CONNECTED) {
+        return -ENOTCONN;
+    }
+
+    if (slot->selected_physical_layout_handle == 0) {
+        // It appears that sometimes the peripheral is considered connected
+        // before the GATT characteristics have been discovered. If this is
+        // the case, the selected_physical_layout_handle will not yet be set.
+        return -EAGAIN;
+    }
+
+    if (bt_conn_get_security(slot->conn) < BT_SECURITY_L2) {
+        return -EAGAIN;
+    }
+
+    int err = bt_gatt_write_without_response(slot->conn, slot->selected_physical_layout_handle,
+                                             &layout_idx, sizeof(layout_idx), true);
+
+    if (err < 0) {
+        LOG_ERR("Failed to write physical layout index to peripheral (err %d)", err);
+    }
+
+    return err;
+}
+
+static void update_peripherals_selected_physical_layout(struct k_work *_work) {
+    uint8_t layout_idx = zmk_physical_layouts_get_selected();
+    for (int i = 0; i < ZMK_SPLIT_BLE_PERIPHERAL_COUNT; i++) {
+        if (peripherals[i].state != PERIPHERAL_SLOT_STATE_CONNECTED) {
+            continue;
+        }
+
+        update_peripheral_selected_layout(&peripherals[i], layout_idx);
+    }
+}
+
+K_WORK_DEFINE(update_peripherals_selected_layouts_work,
+              update_peripherals_selected_physical_layout);
+
 static uint8_t split_central_chrc_discovery_func(struct bt_conn *conn,
                                                  const struct bt_gatt_attr *attr,
                                                  struct bt_gatt_discover_params *params) {
@@ -442,6 +485,11 @@ static uint8_t split_central_chrc_discovery_func(struct bt_conn *conn,
         slot->discover_params.uuid = NULL;
         slot->discover_params.start_handle = attr->handle + 2;
         slot->run_behavior_handle = bt_gatt_attr_value_handle(attr);
+    } else if (!bt_uuid_cmp(((struct bt_gatt_chrc *)attr->user_data)->uuid,
+                            BT_UUID_DECLARE_128(ZMK_SPLIT_BT_SELECT_PHYS_LAYOUT_UUID))) {
+        LOG_DBG("Found select physical layout handle");
+        slot->selected_physical_layout_handle = bt_gatt_attr_value_handle(attr);
+        k_work_submit(&update_peripherals_selected_layouts_work);
 #if IS_ENABLED(CONFIG_ZMK_SPLIT_PERIPHERAL_HID_INDICATORS)
     } else if (!bt_uuid_cmp(((struct bt_gatt_chrc *)attr->user_data)->uuid,
                             BT_UUID_DECLARE_128(ZMK_SPLIT_BT_UPDATE_HID_INDICATORS_UUID))) {
@@ -467,7 +515,8 @@ static uint8_t split_central_chrc_discovery_func(struct bt_conn *conn,
 #endif /* IS_ENABLED(CONFIG_ZMK_SPLIT_BLE_CENTRAL_BATTERY_LEVEL_FETCHING) */
     }
 
-    bool subscribed = slot->run_behavior_handle && slot->subscribe_params.value_handle;
+    bool subscribed = slot->run_behavior_handle && slot->subscribe_params.value_handle &&
+                      slot->selected_physical_layout_handle;
 
 #if ZMK_KEYMAP_HAS_SENSORS
     subscribed = subscribed && slot->sensor_subscribe_params.value_handle;
@@ -739,9 +788,30 @@ static void split_central_disconnected(struct bt_conn *conn, uint8_t reason) {
     start_scanning();
 }
 
+static void split_central_security_changed(struct bt_conn *conn, bt_security_t level,
+                                           enum bt_security_err err) {
+    struct peripheral_slot *slot = peripheral_slot_for_conn(conn);
+    if (!slot || !slot->selected_physical_layout_handle) {
+        return;
+    }
+
+    if (err > 0) {
+        LOG_DBG("Skipping updating the physical layout for peripheral with security error");
+        return;
+    }
+
+    if (level < BT_SECURITY_L2) {
+        LOG_DBG("Skipping updating the physical layout for peripheral with insufficient security");
+        return;
+    }
+
+    k_work_submit(&update_peripherals_selected_layouts_work);
+}
+
 static struct bt_conn_cb conn_callbacks = {
     .connected = split_central_connected,
     .disconnected = split_central_disconnected,
+    .security_changed = split_central_security_changed,
 };
 
 K_THREAD_STACK_DEFINE(split_central_split_run_q_stack,
@@ -897,3 +967,13 @@ static int zmk_split_bt_central_init(void) {
 }
 
 SYS_INIT(zmk_split_bt_central_init, APPLICATION, CONFIG_ZMK_BLE_INIT_PRIORITY);
+
+static int zmk_split_bt_central_listener_cb(const zmk_event_t *eh) {
+    if (as_zmk_physical_layout_selection_changed(eh)) {
+        k_work_submit(&update_peripherals_selected_layouts_work);
+    }
+    return ZMK_EV_EVENT_BUBBLE;
+}
+
+ZMK_LISTENER(zmk_split_bt_central, zmk_split_bt_central_listener_cb);
+ZMK_SUBSCRIPTION(zmk_split_bt_central, zmk_physical_layout_selection_changed);
\ No newline at end of file
diff --git a/app/src/split/bluetooth/service.c b/app/src/split/bluetooth/service.c
index 505eb363cd8..9529d51613c 100644
--- a/app/src/split/bluetooth/service.c
+++ b/app/src/split/bluetooth/service.c
@@ -19,6 +19,7 @@ LOG_MODULE_DECLARE(zmk, CONFIG_ZMK_LOG_LEVEL);
 #include <drivers/behavior.h>
 #include <zmk/behavior.h>
 #include <zmk/matrix.h>
+#include <zmk/physical_layouts.h>
 #include <zmk/split/bluetooth/uuid.h>
 #include <zmk/split/bluetooth/service.h>
 
@@ -138,6 +139,42 @@ static ssize_t split_svc_update_indicators(struct bt_conn *conn, const struct bt
 
 #endif // IS_ENABLED(CONFIG_ZMK_SPLIT_PERIPHERAL_HID_INDICATORS)
 
+static uint8_t selected_phys_layout = 0;
+
+static void split_svc_select_phys_layout_callback(struct k_work *work) {
+    LOG_DBG("Selecting physical layout after GATT write of %d", selected_phys_layout);
+    zmk_physical_layouts_select(selected_phys_layout);
+}
+
+static K_WORK_DEFINE(split_svc_select_phys_layout_work, split_svc_select_phys_layout_callback);
+
+static ssize_t split_svc_select_phys_layout(struct bt_conn *conn, const struct bt_gatt_attr *attr,
+                                            const void *buf, uint16_t len, uint16_t offset,
+                                            uint8_t flags) {
+    if (offset + len > sizeof(uint8_t) || len == 0) {
+        return BT_GATT_ERR(BT_ATT_ERR_INVALID_OFFSET);
+    }
+
+    selected_phys_layout = *(uint8_t *)buf;
+
+    k_work_submit(&split_svc_select_phys_layout_work);
+
+    return len;
+}
+
+static ssize_t split_svc_get_selected_phys_layout(struct bt_conn *conn,
+                                                  const struct bt_gatt_attr *attrs, void *buf,
+                                                  uint16_t len, uint16_t offset) {
+    int selected_ret = zmk_physical_layouts_get_selected();
+    if (selected_ret < 0) {
+        return BT_GATT_ERR(BT_ATT_ERR_VALUE_NOT_ALLOWED);
+    }
+
+    uint8_t selected = (uint8_t)selected_ret;
+
+    return bt_gatt_attr_read(conn, attrs, buf, len, offset, &selected, sizeof(selected));
+}
+
 BT_GATT_SERVICE_DEFINE(
     split_svc, BT_GATT_PRIMARY_SERVICE(BT_UUID_DECLARE_128(ZMK_SPLIT_BT_SERVICE_UUID)),
     BT_GATT_CHARACTERISTIC(BT_UUID_DECLARE_128(ZMK_SPLIT_BT_CHAR_POSITION_STATE_UUID),
@@ -160,7 +197,11 @@ BT_GATT_SERVICE_DEFINE(
                            BT_GATT_CHRC_WRITE_WITHOUT_RESP, BT_GATT_PERM_WRITE_ENCRYPT, NULL,
                            split_svc_update_indicators, NULL),
 #endif // IS_ENABLED(CONFIG_ZMK_SPLIT_PERIPHERAL_HID_INDICATORS)
-);
+    BT_GATT_CHARACTERISTIC(BT_UUID_DECLARE_128(ZMK_SPLIT_BT_SELECT_PHYS_LAYOUT_UUID),
+                           BT_GATT_CHRC_WRITE | BT_GATT_CHRC_READ,
+                           BT_GATT_PERM_WRITE_ENCRYPT | BT_GATT_PERM_READ_ENCRYPT,
+                           split_svc_get_selected_phys_layout, split_svc_select_phys_layout,
+                           NULL), );
 
 K_THREAD_STACK_DEFINE(service_q_stack, CONFIG_ZMK_SPLIT_BLE_PERIPHERAL_STACK_SIZE);