1 files changed, 23 insertions, 14 deletions

diff --git a/app/src/display/widgets/wpm_status.c b/app/src/display/widgets/wpm_status.c
index 41ee368..bee3413 100644
--- a/app/src/display/widgets/wpm_status.c
+++ b/app/src/display/widgets/wpm_status.c
@@ -7,6 +7,7 @@
 #include <logging/log.h>
 LOG_MODULE_DECLARE(zmk, CONFIG_ZMK_LOG_LEVEL);
 
+#include <zmk/display.h>
 #include <zmk/display/widgets/wpm_status.h>
 #include <zmk/events/wpm_state_changed.h>
 #include <zmk/event_manager.h>
@@ -18,6 +19,10 @@ static lv_style_t label_style;
 
 static bool style_initialized = false;
 
+struct wpm_status_state {
+    uint8_t wpm;
+};
+
 void wpm_status_init() {
     if (style_initialized) {
         return;
@@ -31,37 +36,41 @@ void wpm_status_init() {
     lv_style_set_text_line_space(&label_style, LV_STATE_DEFAULT, 1);
 }
 
-void set_wpm_symbol(lv_obj_t *label, int wpm) {
+struct wpm_status_state wpm_status_get_state(const zmk_event_t *eh) {
+    return (struct wpm_status_state){.wpm = zmk_wpm_get_state()};
+};
+
+void set_wpm_symbol(lv_obj_t *label, struct wpm_status_state state) {
     char text[4] = {};
 
-    LOG_DBG("WPM changed to %i", wpm);
-    sprintf(text, "%i ", wpm);
+    LOG_DBG("WPM changed to %i", state.wpm);
+    snprintf(text, sizeof(text), "%i ", state.wpm);
 
     lv_label_set_text(label, text);
 }
 
+void wpm_status_update_cb(struct wpm_status_state state) {
+    struct zmk_widget_wpm_status *widget;
+    SYS_SLIST_FOR_EACH_CONTAINER(&widgets, widget, node) { set_wpm_symbol(widget->obj, state); }
+}
+
+ZMK_DISPLAY_WIDGET_LISTENER(widget_wpm_status, struct wpm_status_state, wpm_status_update_cb,
+                            wpm_status_get_state)
+ZMK_SUBSCRIPTION(widget_wpm_status, zmk_wpm_state_changed);
+
 int zmk_widget_wpm_status_init(struct zmk_widget_wpm_status *widget, lv_obj_t *parent) {
     wpm_status_init();
+
     widget->obj = lv_label_create(parent, NULL);
     lv_obj_add_style(widget->obj, LV_LABEL_PART_MAIN, &label_style);
     lv_label_set_align(widget->obj, LV_LABEL_ALIGN_RIGHT);
 
     lv_obj_set_size(widget->obj, 40, 15);
-    set_wpm_symbol(widget->obj, 0);
 
     sys_slist_append(&widgets, &widget->node);
 
+    widget_wpm_status_init();
     return 0;
 }
 
 lv_obj_t *zmk_widget_wpm_status_obj(struct zmk_widget_wpm_status *widget) { return widget->obj; }
-
-int wpm_status_listener(const zmk_event_t *eh) {
-    struct zmk_wpm_state_changed *ev = as_zmk_wpm_state_changed(eh);
-    struct zmk_widget_wpm_status *widget;
-    SYS_SLIST_FOR_EACH_CONTAINER(&widgets, widget, node) { set_wpm_symbol(widget->obj, ev->state); }
-    return 0;
-}
-
-ZMK_LISTENER(widget_wpm_status, wpm_status_listener)
-ZMK_SUBSCRIPTION(widget_wpm_status, zmk_wpm_state_changed);
\ No newline at end of file