diff --git a/stm32-modules/flex-stacker/firmware/CMakeLists.txt b/stm32-modules/flex-stacker/firmware/CMakeLists.txt
index af0368367..4b5c63034 100644
--- a/stm32-modules/flex-stacker/firmware/CMakeLists.txt
+++ b/stm32-modules/flex-stacker/firmware/CMakeLists.txt
@@ -21,7 +21,9 @@ set(${TARGET_MODULE_NAME}_FW_LINTABLE_SRCS
     ${SYSTEM_DIR}/freertos_idle_timer_task.cpp
     ${SYSTEM_DIR}/freertos_system_task.cpp
     ${SYSTEM_DIR}/system_policy.cpp
+    ${SYSTEM_DIR}/i2c_comms.cpp
     ${UI_DIR}/freertos_ui_task.cpp
+    ${UI_DIR}/ui_policy.cpp
     ${MOTOR_CONTROL_DIR}/freertos_motor_task.cpp
     ${MOTOR_CONTROL_DIR}/freertos_motor_driver_task.cpp
     ${MOTOR_CONTROL_DIR}/motor_driver_policy.cpp
@@ -36,6 +38,7 @@ set(${TARGET_MODULE_NAME}_FW_NONLINTABLE_SRCS
     ${SYSTEM_DIR}/hal_tick.c
     ${SYSTEM_DIR}/system_serial_number.c
     ${SYSTEM_DIR}/system_hardware.c
+    ${SYSTEM_DIR}/i2c_hardware.c
     ${UI_DIR}/ui_hardware.c
     ${COMMS_DIR}/usbd_conf.c
     ${COMMS_DIR}/usbd_desc.c
diff --git a/stm32-modules/flex-stacker/firmware/main.cpp b/stm32-modules/flex-stacker/firmware/main.cpp
index bb09f4f91..c44959093 100644
--- a/stm32-modules/flex-stacker/firmware/main.cpp
+++ b/stm32-modules/flex-stacker/firmware/main.cpp
@@ -1,12 +1,15 @@
 #include "FreeRTOS.h"
+#include "task.h"
+
+#include "ot_utils/freertos/freertos_task.hpp"
 #include "firmware/firmware_tasks.hpp"
 #include "firmware/freertos_tasks.hpp"
 #include "firmware/motor_hardware.h"
+#include "firmware/i2c_comms.hpp"
+#include "firmware/i2c_hardware.h"
 #include "firmware/system_stm32g4xx.h"
 #include "flex-stacker/messages.hpp"
-#include "ot_utils/freertos/freertos_task.hpp"
 #include "systemwide.h"
-#include "task.h"
 
 #pragma GCC diagnostic push
 // NOLINTNEXTLINE(clang-diagnostic-unknown-warning-option)
@@ -15,6 +18,8 @@
 #pragma GCC diagnostic pop
 
 using EntryPoint = std::function<void(tasks::FirmwareTasks::QueueAggregator *)>;
+using EntryPointUI = std::function<void(
+    tasks::FirmwareTasks::QueueAggregator *, i2c::hardware::I2C *)>;
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
 static auto motor_driver_task_entry = EntryPoint(motor_driver_task::run);
@@ -23,7 +28,7 @@ static auto motor_driver_task_entry = EntryPoint(motor_driver_task::run);
 static auto motor_task_entry = EntryPoint(motor_control_task::run);
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-static auto ui_task_entry = EntryPoint(ui_control_task::run);
+static auto ui_task_entry = EntryPointUI(ui_control_task::run);
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
 static auto host_comms_entry = EntryPoint(host_comms_control_task::run);
@@ -51,7 +56,7 @@ static auto host_comms_task =
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
 static auto ui_task =
-    ot_utils::freertos_task::FreeRTOSTask<tasks::UI_STACK_SIZE, EntryPoint>(
+    ot_utils::freertos_task::FreeRTOSTask<tasks::UI_STACK_SIZE, EntryPointUI>(
         ui_task_entry);
 
 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
@@ -71,18 +76,22 @@ extern "C" void HAL_GPIO_EXTI_Callback(uint16_t GPIO_Pin) {
     }
 }
 
+static auto i2c2_comms = i2c::hardware::I2C();
+static auto i2c3_comms = i2c::hardware::I2C();
+static auto i2c_handles = I2CHandlerStruct{};
+
 auto main() -> int {
     HardwareInit();
 
-    system_task.start(tasks::SYSTEM_TASK_PRIORITY, "System", &aggregator);
+    i2c_hardware_init(&i2c_handles);
+    i2c2_comms.set_handle(i2c_handles.i2c2, I2C_BUS_2);
+    i2c3_comms.set_handle(i2c_handles.i2c3, I2C_BUS_3);
 
-    driver_task.start(tasks::MOTOR_DRIVER_TASK_PRIORITY, "Motor Driver",
-                      &aggregator);
+    system_task.start(tasks::SYSTEM_TASK_PRIORITY, "System", &aggregator);
+    driver_task.start(tasks::MOTOR_DRIVER_TASK_PRIORITY, "Motor Driver", &aggregator);
     motor_task.start(tasks::MOTOR_TASK_PRIORITY, "Motor", &aggregator);
-
     host_comms_task.start(tasks::COMMS_TASK_PRIORITY, "Comms", &aggregator);
-
-    ui_task.start(tasks::UI_TASK_PRIORITY, "UI", &aggregator);
+    ui_task.start(tasks::UI_TASK_PRIORITY, "UI", &aggregator, &i2c2_comms);
 
     vTaskStartScheduler();
     return 0;
diff --git a/stm32-modules/flex-stacker/firmware/system/i2c_comms.cpp b/stm32-modules/flex-stacker/firmware/system/i2c_comms.cpp
new file mode 100644
index 000000000..65dee66eb
--- /dev/null
+++ b/stm32-modules/flex-stacker/firmware/system/i2c_comms.cpp
@@ -0,0 +1,27 @@
+#include <stdint.h>
+
+#include "firmware/i2c_comms.hpp"
+#include "firmware/i2c_hardware.h"
+#include "systemwide.h"
+
+using namespace i2c::hardware;
+
+auto I2C::set_handle(HAL_I2C_HANDLE i2c_handle, I2C_BUS i2c_bus) -> void {
+    this->handle = i2c_handle;
+    this->bus = i2c_bus;
+    i2c_register_handle(this->handle, this->bus);
+}
+
+auto I2C::i2c_write(uint16_t dev_addr, uint16_t reg, uint8_t* data,
+                    uint16_t size) -> RxTxReturn {
+    MessageT resp{0};
+    auto ret = hal_i2c_write(bus, dev_addr, reg, data, size, TIMEOUT);
+    return RxTxReturn(ret, resp);
+}
+
+auto I2C::i2c_read(uint16_t dev_addr, uint16_t reg, uint16_t size)
+    -> RxTxReturn {
+    MessageT resp{0};
+    auto ret = hal_i2c_read(bus, dev_addr, reg, resp.data(), size, TIMEOUT);
+    return RxTxReturn(ret, resp);
+}
diff --git a/stm32-modules/flex-stacker/firmware/system/i2c_hardware.c b/stm32-modules/flex-stacker/firmware/system/i2c_hardware.c
new file mode 100644
index 000000000..8a44b2b19
--- /dev/null
+++ b/stm32-modules/flex-stacker/firmware/system/i2c_hardware.c
@@ -0,0 +1,305 @@
+#include <stdint.h>
+
+#include "stm32g4xx_hal.h"
+#include "stm32g4xx_hal_conf.h"
+#include "stm32g4xx_hal_gpio.h"
+#include "stm32g4xx_hal_def.h"
+#include "stm32g4xx_it.h"
+#include "FreeRTOS.h"
+#include "task.h"
+
+#include "firmware/i2c_hardware.h"
+
+#define MAX_I2C_HANDLES (3)
+#define NO_HANDLE_ERROR (255)
+/** Size of register address: 1 byte.*/
+#define REGISTER_ADDR_LEN (1)
+
+static I2C_HandleTypeDef hi2c2;
+static I2C_HandleTypeDef hi2c3;
+
+typedef struct {
+    I2C_BUS bus;
+    I2C_HandleTypeDef *i2c_handle;
+    TaskHandle_t task_to_notify;
+    bool should_retry;
+} NotificationHandle_t;
+
+static NotificationHandle_t _notification_handles[MAX_I2C_HANDLES];
+
+static bool _initialized = false;
+
+HAL_I2C_HANDLE MX_I2C2_Init()
+{
+    hi2c2.State = HAL_I2C_STATE_RESET;
+    hi2c2.Instance = I2C2;
+    hi2c2.Init.Timing = 0x10C0ECFF;
+    hi2c2.Init.OwnAddress1 = 0;
+    hi2c2.Init.AddressingMode = I2C_ADDRESSINGMODE_7BIT;
+    hi2c2.Init.DualAddressMode = I2C_DUALADDRESS_DISABLE;
+    hi2c2.Init.OwnAddress2 = 0;
+    hi2c2.Init.OwnAddress2Masks = I2C_OA2_NOMASK;
+    hi2c2.Init.GeneralCallMode = I2C_GENERALCALL_DISABLE;
+    hi2c2.Init.NoStretchMode = I2C_NOSTRETCH_DISABLE;
+    if (HAL_I2C_Init(&hi2c2) != HAL_OK)
+    {
+        Error_Handler();
+    }
+    /** Configure Analogue filter
+    */
+    if (HAL_I2CEx_ConfigAnalogFilter(&hi2c2, I2C_ANALOGFILTER_ENABLE) != HAL_OK)
+    {
+        Error_Handler();
+    }
+    /** Configure Digital filter
+    */
+    if (HAL_I2CEx_ConfigDigitalFilter(&hi2c2, 0) != HAL_OK)
+    {
+        Error_Handler();
+    }
+
+    /** I2C Fast mode Plus enable */
+    __HAL_SYSCFG_FASTMODEPLUS_ENABLE(I2C_FASTMODEPLUS_I2C2);
+
+    return &hi2c2;
+}
+
+HAL_I2C_HANDLE MX_I2C3_Init() {
+    hi2c3.Instance = I2C3;
+    hi2c3.Init.Timing = 0x10C0ECFF;
+    hi2c3.Init.OwnAddress1 = 0;
+    hi2c3.Init.AddressingMode = I2C_ADDRESSINGMODE_7BIT;
+    hi2c3.Init.DualAddressMode = I2C_DUALADDRESS_DISABLE;
+    hi2c3.Init.OwnAddress2 = 0;
+    hi2c3.Init.OwnAddress2Masks = I2C_OA2_NOMASK;
+    hi2c3.Init.GeneralCallMode = I2C_GENERALCALL_DISABLE;
+    hi2c3.Init.NoStretchMode = I2C_NOSTRETCH_DISABLE;
+    if (HAL_I2C_Init(&hi2c3) != HAL_OK) {
+        Error_Handler();
+    }
+    /** Configure Analogue filter
+     */
+    if (HAL_I2CEx_ConfigAnalogFilter(&hi2c3, I2C_ANALOGFILTER_ENABLE) !=
+        HAL_OK) {
+        Error_Handler();
+    }
+    /** Configure Digital filter
+     */
+    if (HAL_I2CEx_ConfigDigitalFilter(&hi2c3, 0) != HAL_OK) {
+        Error_Handler();
+    }
+
+    /** I2C Fast mode Plus enable */
+    __HAL_SYSCFG_FASTMODEPLUS_ENABLE(I2C_FASTMODEPLUS_I2C3);
+
+    return &hi2c3;
+}
+
+void i2c_hardware_init(I2CHandlerStruct* i2c_handles) {
+    HAL_I2C_HANDLE i2c2 = MX_I2C2_Init();
+    HAL_I2C_HANDLE i2c3 = MX_I2C3_Init();
+    i2c_handles->i2c2 = i2c2;
+    i2c_handles->i2c3 = i2c3;
+}
+
+/**
+ * @brief Get the notification handle based on the I2C Bus.
+ * Returns NULL if the handle is not found.
+ */
+static NotificationHandle_t* lookup_handle(I2C_BUS bus) {
+    for(size_t i = 0; i < MAX_I2C_HANDLES; ++i) {
+        I2C_BUS notif_bus = _notification_handles[i].bus;
+        if(notif_bus != NO_BUS && notif_bus == bus) {
+            return &_notification_handles[i];
+        }
+    }
+    return NULL;
+}
+
+/**
+ * @brief Get the notification handle based on the HAL I2C struct.
+ * Returns NULL if the handle is not found.
+ */
+static NotificationHandle_t* get_handle(I2C_HandleTypeDef *i2c_handle) {
+    for(size_t i = 0; i < MAX_I2C_HANDLES; ++i) {
+        if(_notification_handles[i].i2c_handle == i2c_handle) {
+            return &_notification_handles[i];
+        }
+    }
+    return NULL;
+}
+
+static void initialize_notification_handles() {
+    if(!_initialized) {
+        for(size_t i = 0; i < MAX_I2C_HANDLES; ++i) {
+            _notification_handles[i].i2c_handle = NULL;
+            _notification_handles[i].task_to_notify = NULL;
+            _notification_handles[i].should_retry = false;
+            _notification_handles[i].bus = NO_BUS;
+        }
+        _initialized = true;
+    }
+}
+
+/**
+ * @brief Common handler for all I2C callbacks.
+ */
+static void handle_i2c_callback(I2C_HandleTypeDef *i2c_handle, bool error) {
+    BaseType_t xHigherPriorityTaskWoken = pdFALSE;
+    NotificationHandle_t *instance = get_handle(i2c_handle);
+    if(instance == NULL) {
+        return;
+    }
+    if(instance->task_to_notify == NULL) {
+        return;
+    }
+
+    vTaskNotifyGiveFromISR(instance->task_to_notify, 
+                            &xHigherPriorityTaskWoken);
+    instance->task_to_notify = NULL;
+    instance->should_retry = error;
+    portYIELD_FROM_ISR( xHigherPriorityTaskWoken );
+}
+
+bool i2c_register_handle(HAL_I2C_HANDLE handle, I2C_BUS bus) {
+    initialize_notification_handles();
+
+    I2C_HandleTypeDef* i2c_handle = (I2C_HandleTypeDef*)handle;
+    NotificationHandle_t *notif_handle = lookup_handle(bus);
+    if(notif_handle != NULL) {
+        // Already registered
+        return true;
+    }
+
+    // Now find an empty slot
+    notif_handle = get_handle(NULL);
+    if(notif_handle == NULL) {
+        // No empty slots
+        return false;
+    }
+    notif_handle->i2c_handle = i2c_handle;
+    notif_handle->bus = bus;
+    return true;
+}
+
+uint8_t hal_i2c_comms_ready(HAL_I2C_HANDLE handle, uint16_t dev_address, uint8_t tries, uint32_t timeout) {
+    HAL_StatusTypeDef status = HAL_OK;
+    status = HAL_I2C_IsDeviceReady(handle, dev_address, tries, timeout);
+    return (uint8_t)status;
+}
+
+/**
+ * Wrapper around HAL_I2C_Mem_Write
+ */
+uint8_t hal_i2c_write(I2C_BUS bus, uint16_t DevAddress, uint16_t reg, uint8_t *data, uint16_t size, uint32_t timeout) {
+    uint32_t notification_val = 0;
+    NotificationHandle_t *notification_handle = lookup_handle(bus);
+    I2C_HandleTypeDef* i2c_handle = (I2C_HandleTypeDef*)notification_handle->i2c_handle;
+
+    if(notification_handle == NULL) {
+        return NO_HANDLE_ERROR;
+    }
+
+    // Make sure the device is online
+    HAL_StatusTypeDef dev_status = HAL_OK;
+    dev_status = hal_i2c_comms_ready(i2c_handle, DevAddress, 3, timeout);
+    if (dev_status != HAL_OK) return dev_status;
+
+    uint32_t tickstart = HAL_GetTick();
+    HAL_StatusTypeDef tx_result = HAL_OK;
+    do {
+        notification_handle->task_to_notify = xTaskGetCurrentTaskHandle();
+        tx_result = HAL_I2C_Mem_Write_IT(i2c_handle,
+                            DevAddress, reg, REGISTER_ADDR_LEN, data, size);
+        notification_val = ulTaskNotifyTake(pdTRUE, timeout); // Wait for callback
+        if (notification_handle->should_retry) {
+            tx_result = HAL_BUSY;
+        }
+        if(notification_val != 1) {
+            // Interrupt never fired
+            tx_result = HAL_TIMEOUT;
+        }
+        if (tx_result == HAL_OK) {
+            break;
+        }
+    } while ((HAL_GetTick() - tickstart) < timeout);
+    return (uint8_t)tx_result;
+}
+
+/**
+ * Wrapper around HAL_I2C_Mem_Read
+ */
+uint8_t hal_i2c_read(I2C_BUS bus, uint16_t DevAddress, uint16_t reg, uint8_t *data, uint16_t size, uint32_t timeout) {
+    uint32_t notification_val = 0;
+    NotificationHandle_t *notification_handle = lookup_handle(bus);
+    I2C_HandleTypeDef* i2c_handle = (I2C_HandleTypeDef*)notification_handle->i2c_handle;
+
+    if(notification_handle == NULL) {
+        return NO_HANDLE_ERROR;
+    }
+
+    // Make sure the device is online
+    HAL_StatusTypeDef dev_status = HAL_OK;
+    dev_status = hal_i2c_comms_ready(i2c_handle, DevAddress, 3, timeout);
+    if (dev_status != HAL_OK) return dev_status;
+
+    uint32_t tickstart = HAL_GetTick();
+    HAL_StatusTypeDef rx_result = HAL_OK;
+    do {
+        notification_handle->task_to_notify = xTaskGetCurrentTaskHandle();
+        rx_result = HAL_I2C_Mem_Read_IT(i2c_handle,
+                            DevAddress, reg, REGISTER_ADDR_LEN, data, size);
+        notification_val = ulTaskNotifyTake(pdTRUE, timeout); // Wait for callback
+        if (notification_handle->should_retry) {
+            rx_result = HAL_BUSY;
+        }
+        if(notification_val != 1) {
+            // Interrupt never fired
+            rx_result = HAL_TIMEOUT;
+        }
+        if (rx_result == HAL_OK) {
+            break;
+        }
+    } while ((HAL_GetTick() - tickstart) < timeout);
+    return (uint8_t)rx_result;
+}
+
+void HAL_I2C_MemTxCpltCallback(I2C_HandleTypeDef *i2c_handle) {
+    handle_i2c_callback(i2c_handle, false);
+}
+
+void HAL_I2C_MemRxCpltCallback(I2C_HandleTypeDef *i2c_handle) {
+    handle_i2c_callback(i2c_handle, false);
+}
+
+void HAL_I2C_MasterTxCpltCallback(I2C_HandleTypeDef *i2c_handle) {
+    handle_i2c_callback(i2c_handle, false);
+}
+
+void HAL_I2C_MasterRxCpltCallback(I2C_HandleTypeDef *i2c_handle) {
+    handle_i2c_callback(i2c_handle, false);
+}
+
+void HAL_I2C_ErrorCallback(I2C_HandleTypeDef *i2c_handle) {
+    handle_i2c_callback(i2c_handle, true);
+}
+
+void I2C2_EV_IRQHandler(void)
+{
+    HAL_I2C_EV_IRQHandler(&hi2c2);
+}
+
+void I2C2_ER_IRQHandler(void)
+{
+    HAL_I2C_ER_IRQHandler(&hi2c2);
+}
+
+void I2C3_EV_IRQHandler(void)
+{
+    HAL_I2C_EV_IRQHandler(&hi2c3);
+}
+
+void I2C3_ER_IRQHandler(void)
+{
+    HAL_I2C_ER_IRQHandler(&hi2c3);
+}
diff --git a/stm32-modules/flex-stacker/firmware/system/main.h b/stm32-modules/flex-stacker/firmware/system/main.h
index 5d5b0f1f3..20da410de 100644
--- a/stm32-modules/flex-stacker/firmware/system/main.h
+++ b/stm32-modules/flex-stacker/firmware/system/main.h
@@ -72,10 +72,14 @@ void Error_Handler(void);
 #define TOF_I2C3_SCL_GPIO_Port GPIOC
 #define TOF_I2C3_SDA_Pin GPIO_PIN_9
 #define TOF_I2C3_SDA_GPIO_Port GPIOC
+
 #define EEPROM_I2C2_SDA_Pin GPIO_PIN_8
 #define EEPROM_I2C2_SDA_GPIO_Port GPIOA
 #define EEPOM_I2C2_SCL_Pin GPIO_PIN_9
 #define EEPOM_I2C2_SCL_GPIO_Port GPIOA
+#define EEPROM_WP_PIN GPIO_PIN_10
+#define EEPROM_WP_PORT GPIOA
+
 #define nSTATUS_LED_Pin GPIO_PIN_10
 #define nSTATUS_LED_GPIO_Port GPIOC
 #define nLW_RELEASED_Pin GPIO_PIN_11
diff --git a/stm32-modules/flex-stacker/firmware/system/stm32g4xx_hal_msp.c b/stm32-modules/flex-stacker/firmware/system/stm32g4xx_hal_msp.c
index bec6e11c1..1333aa156 100644
--- a/stm32-modules/flex-stacker/firmware/system/stm32g4xx_hal_msp.c
+++ b/stm32-modules/flex-stacker/firmware/system/stm32g4xx_hal_msp.c
@@ -73,18 +73,8 @@ void HAL_MspDeInit(void)
 void HAL_I2C_MspInit(I2C_HandleTypeDef* hi2c)
 {
     GPIO_InitTypeDef GPIO_InitStruct = {0};
-    RCC_PeriphCLKInitTypeDef PeriphClkInit = {0};
     if(hi2c->Instance==I2C2)
     {
-        /** Initializes the peripherals clocks
-         */
-        PeriphClkInit.PeriphClockSelection = RCC_PERIPHCLK_I2C2;
-        PeriphClkInit.I2c2ClockSelection = RCC_I2C2CLKSOURCE_PCLK1;
-        if (HAL_RCCEx_PeriphCLKConfig(&PeriphClkInit) != HAL_OK)
-        {
-            Error_Handler();
-        }
-
         __HAL_RCC_GPIOA_CLK_ENABLE();
         /**I2C2 GPIO Configuration
         PA8     ------> I2C2_SDA
@@ -99,18 +89,14 @@ void HAL_I2C_MspInit(I2C_HandleTypeDef* hi2c)
 
         /* Peripheral clock enable */
         __HAL_RCC_I2C2_CLK_ENABLE();
+
+        HAL_NVIC_SetPriority(I2C2_EV_IRQn, 7, 0);
+        HAL_NVIC_SetPriority(I2C2_ER_IRQn, 7, 0);
+        HAL_NVIC_EnableIRQ(I2C2_EV_IRQn);
+        HAL_NVIC_EnableIRQ(I2C2_ER_IRQn);
     }
     else if(hi2c->Instance==I2C3)
     {
-        /** Initializes the peripherals clocks
-         */
-        PeriphClkInit.PeriphClockSelection = RCC_PERIPHCLK_I2C3;
-        PeriphClkInit.I2c3ClockSelection = RCC_I2C3CLKSOURCE_PCLK1;
-        if (HAL_RCCEx_PeriphCLKConfig(&PeriphClkInit) != HAL_OK)
-        {
-            Error_Handler();
-        }
-
         __HAL_RCC_GPIOC_CLK_ENABLE();
         /**I2C3 GPIO Configuration
         PC8     ------> I2C3_SCL
@@ -125,8 +111,12 @@ void HAL_I2C_MspInit(I2C_HandleTypeDef* hi2c)
 
         /* Peripheral clock enable */
         __HAL_RCC_I2C3_CLK_ENABLE();
-    }
 
+        HAL_NVIC_SetPriority(I2C3_EV_IRQn, 7, 0);
+        HAL_NVIC_SetPriority(I2C3_ER_IRQn, 7, 0);
+        HAL_NVIC_EnableIRQ(I2C3_EV_IRQn);
+        HAL_NVIC_EnableIRQ(I2C3_ER_IRQn);
+    }
 }
 
 /**
diff --git a/stm32-modules/flex-stacker/firmware/ui/freertos_ui_task.cpp b/stm32-modules/flex-stacker/firmware/ui/freertos_ui_task.cpp
index a69376669..212a69e56 100644
--- a/stm32-modules/flex-stacker/firmware/ui/freertos_ui_task.cpp
+++ b/stm32-modules/flex-stacker/firmware/ui/freertos_ui_task.cpp
@@ -1,17 +1,50 @@
+#include "FreeRTOS.h"
+
 #include "firmware/freertos_tasks.hpp"
 #include "firmware/ui_hardware.h"
+#include "firmware/i2c_comms.hpp"
+#include "firmware/ui_policy.hpp"
+#include "ot_utils/freertos/freertos_timer.hpp"
+#include "flex-stacker/ui_task.hpp"
 
 static constexpr uint32_t HALF_SECOND = 500;
 
 namespace ui_control_task {
+using namespace ui_policy;
+
+enum class Notifications : uint8_t {
+    INCOMING_MESSAGE = 1,
+};
+
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+static tasks::FirmwareTasks::UIQueue
+    // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+    _queue(static_cast<uint8_t>(Notifications::INCOMING_MESSAGE), "UI Queue");
+
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+static auto _top_task = ui_task::UITask(_queue, nullptr);
+
+
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+static auto _timer_callback = [ObjectPtr = &_top_task] {
+    ObjectPtr->update_callback();
+};
+
+// This timer provides the backing timing for the UI task
+// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
+static auto _ui_timer = ot_utils::freertos_timer::FreeRTOSTimer(
+    "UI Timer", _timer_callback, decltype(_top_task)::UPDATE_PERIOD_MS);
+
+auto run(tasks::FirmwareTasks::QueueAggregator* aggregator, i2c::hardware::I2C* i2c_comms) -> void {
+    auto* handle = xTaskGetCurrentTaskHandle();
+    _queue.provide_handle(handle);
+    aggregator->register_queue(_queue);
+    _top_task.provide_aggregator(aggregator);
 
-auto run(tasks::FirmwareTasks::QueueAggregator* aggregator) -> void {
-    std::ignore = aggregator;
+    auto policy = UIPolicy();
+    _ui_timer.start();
     while (true) {
-        ui_hardware_set_heartbeat_led(true);
-        vTaskDelay(HALF_SECOND);
-        ui_hardware_set_heartbeat_led(false);
-        vTaskDelay(HALF_SECOND);
+        _top_task.run_once(policy);
     }
 }
 
diff --git a/stm32-modules/flex-stacker/firmware/ui/ui_hardware.c b/stm32-modules/flex-stacker/firmware/ui/ui_hardware.c
index cf57bde44..316a55822 100644
--- a/stm32-modules/flex-stacker/firmware/ui/ui_hardware.c
+++ b/stm32-modules/flex-stacker/firmware/ui/ui_hardware.c
@@ -1,18 +1,10 @@
-#include "main.h"
-
 #include <stdbool.h>
 
 #include "stm32g4xx_hal.h"
-
-
-#define nSTATUS_LED_Pin GPIO_PIN_10
-#define nSTATUS_LED_GPIO_Port GPIOC
+#include "main.h"
 
 void ui_hardware_set_heartbeat_led(bool setting) {
     //NOLINTNEXTLINE(performance-no-int-to-ptr)
     HAL_GPIO_WritePin(nSTATUS_LED_GPIO_Port, nSTATUS_LED_Pin,
         setting ? GPIO_PIN_SET : GPIO_PIN_RESET);
-
-//    HAL_GPIO_WritePin(GPIOC, GPIO_PIN_2,
-//                      setting ? GPIO_PIN_SET : GPIO_PIN_RESET);
 }
diff --git a/stm32-modules/flex-stacker/firmware/ui/ui_policy.cpp b/stm32-modules/flex-stacker/firmware/ui/ui_policy.cpp
new file mode 100644
index 000000000..9062b0aee
--- /dev/null
+++ b/stm32-modules/flex-stacker/firmware/ui/ui_policy.cpp
@@ -0,0 +1,12 @@
+#include "firmware/ui_policy.hpp"
+
+#include "firmware/i2c_hardware.h"
+#include "firmware/ui_hardware.h"
+
+namespace ui_policy {
+
+// NOLINTNEXTLINE(readability-convert-member-functions-to-static)
+auto UIPolicy::set_heartbeat_led(bool value) -> void {
+    ui_hardware_set_heartbeat_led(value);
+}
+};  //ui_policy
diff --git a/stm32-modules/include/flex-stacker/firmware/freertos_tasks.hpp b/stm32-modules/include/flex-stacker/firmware/freertos_tasks.hpp
index 016cd4a51..9deaecabf 100644
--- a/stm32-modules/include/flex-stacker/firmware/freertos_tasks.hpp
+++ b/stm32-modules/include/flex-stacker/firmware/freertos_tasks.hpp
@@ -1,9 +1,11 @@
 #pragma once
 
 #include "FreeRTOS.h"
+#include "task.h"
+
 #include "firmware/firmware_tasks.hpp"
 #include "firmware/freertos_message_queue.hpp"
-#include "task.h"
+#include "firmware/i2c_comms.hpp"
 
 namespace motor_driver_task {
 
@@ -22,7 +24,7 @@ auto run(tasks::FirmwareTasks::QueueAggregator* aggregator) -> void;
 namespace ui_control_task {
 
 // Actual function that runs in the task
-auto run(tasks::FirmwareTasks::QueueAggregator* aggregator) -> void;
+auto run(tasks::FirmwareTasks::QueueAggregator* aggregator, i2c::hardware::I2C* i2c_comms) -> void;
 }  // namespace ui_control_task
 
 namespace host_comms_control_task {
diff --git a/stm32-modules/include/flex-stacker/firmware/hardware_iface.hpp b/stm32-modules/include/flex-stacker/firmware/hardware_iface.hpp
new file mode 100644
index 000000000..52ccc8f29
--- /dev/null
+++ b/stm32-modules/include/flex-stacker/firmware/hardware_iface.hpp
@@ -0,0 +1,33 @@
+#pragma once
+
+#include <stdint.h>
+
+#include <cstdint>
+#include <optional>
+#include <tuple>
+
+using std::size_t;
+static constexpr size_t MESSAGE_LEN = 5;
+using MessageT = std::array<uint8_t, MESSAGE_LEN>;
+
+namespace i2c {
+namespace hardware {
+using RxTxReturn = std::tuple<uint8_t, MessageT>;
+class I2CBase {
+  public:
+    I2CBase() = default;
+    virtual ~I2CBase() = default;
+    I2CBase(const I2CBase&) = default;
+    auto operator=(const I2CBase&) -> I2CBase& = default;
+    I2CBase(I2CBase&&) = default;
+    auto operator=(I2CBase&&) -> I2CBase& = default;
+
+    virtual auto i2c_read(uint16_t dev_addr, uint16_t reg, uint16_t size)
+        -> RxTxReturn;
+    virtual auto i2c_write(uint16_t dev_addr, uint16_t reg, uint8_t* data,
+                           uint16_t size) -> RxTxReturn;
+};
+
+};  // namespace hardware
+
+};  // namespace i2c
diff --git a/stm32-modules/include/flex-stacker/firmware/i2c_comms.hpp b/stm32-modules/include/flex-stacker/firmware/i2c_comms.hpp
new file mode 100644
index 000000000..a16fde215
--- /dev/null
+++ b/stm32-modules/include/flex-stacker/firmware/i2c_comms.hpp
@@ -0,0 +1,35 @@
+#pragma once
+
+#include <stdint.h>
+#include <algorithm>
+#include <cstdint>
+#include <optional>
+
+#include "firmware/hardware_iface.hpp"
+#include "firmware/i2c_hardware.h"
+#include "systemwide.h"
+
+namespace i2c {
+namespace hardware {
+class I2C : public I2CBase {
+  public:
+    explicit I2C() = default;
+    ~I2C() final = default;
+    I2C(const I2C &) = delete;
+    I2C(const I2C &&) = delete;
+    auto operator=(const I2C &) = delete;
+    auto operator=(const I2C &&) = delete;
+
+    auto i2c_read(uint16_t dev_addr, uint16_t reg, uint16_t size) -> RxTxReturn;
+    auto i2c_write(uint16_t dev_addr, uint16_t reg, uint8_t *data,
+                   uint16_t size) -> RxTxReturn;
+    auto set_handle(HAL_I2C_HANDLE i2c_handle, I2C_BUS bus) -> void;
+
+  private:
+    I2C_BUS bus{};
+    HAL_I2C_HANDLE handle = nullptr;
+    // Timeout in ms
+    static constexpr auto TIMEOUT = 1000;
+};
+};  // namespace hardware
+};  // namespace i2c
diff --git a/stm32-modules/include/flex-stacker/firmware/i2c_hardware.h b/stm32-modules/include/flex-stacker/firmware/i2c_hardware.h
new file mode 100644
index 000000000..5bba76019
--- /dev/null
+++ b/stm32-modules/include/flex-stacker/firmware/i2c_hardware.h
@@ -0,0 +1,34 @@
+#pragma once
+
+#include <stdbool.h>
+#include <stdint.h>
+
+#include "systemwide.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif  // __cplusplus
+
+typedef enum I2C_BUS {
+    I2C_BUS_2,
+    I2C_BUS_3,
+    NO_BUS,
+} I2C_BUS;
+
+typedef void *HAL_I2C_HANDLE;
+
+typedef struct HandlerStruct {
+    HAL_I2C_HANDLE i2c2;
+    HAL_I2C_HANDLE i2c3;
+} I2CHandlerStruct;
+
+void i2c_hardware_init(I2CHandlerStruct *i2c_handles);
+bool i2c_register_handle(HAL_I2C_HANDLE handle, I2C_BUS bus);
+uint8_t hal_i2c_write(I2C_BUS bus, uint16_t DevAddress, uint16_t reg,
+                      uint8_t *data, uint16_t size, uint32_t timeout);
+uint8_t hal_i2c_read(I2C_BUS bus, uint16_t DevAddress, uint16_t reg,
+                     uint8_t *data, uint16_t size, uint32_t timeout);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif  // __cplusplus
diff --git a/stm32-modules/include/flex-stacker/firmware/ui_policy.hpp b/stm32-modules/include/flex-stacker/firmware/ui_policy.hpp
new file mode 100644
index 000000000..e67452ceb
--- /dev/null
+++ b/stm32-modules/include/flex-stacker/firmware/ui_policy.hpp
@@ -0,0 +1,23 @@
+#pragma once
+
+#include <array>
+
+#include "firmware/i2c_hardware.h"
+
+namespace ui_policy {
+
+class UIPolicy {
+  public:
+    auto set_heartbeat_led(bool value) -> void;
+    template <size_t Len>
+    auto i2c_write(uint8_t device_address, uint8_t register_address,
+                   std::array<uint8_t, Len> &data) -> bool {
+        auto res = hal_i2c_write(I2C_BUS_2, device_address,
+                                      register_address, data.data(), Len, 1);
+        if (res != 0) {
+            return false;
+        }
+        return res;
+    }
+};
+}; // ui_policy
diff --git a/stm32-modules/include/flex-stacker/flex-stacker/messages.hpp b/stm32-modules/include/flex-stacker/flex-stacker/messages.hpp
index d9e6df612..9e5766905 100644
--- a/stm32-modules/include/flex-stacker/flex-stacker/messages.hpp
+++ b/stm32-modules/include/flex-stacker/flex-stacker/messages.hpp
@@ -272,6 +272,10 @@ struct GetEstopResponse {
     bool triggered;
 };
 
+// This empty message is just used to signal that the UI task should update
+// its outputs
+struct UpdateUIMessage {};
+
 using HostCommsMessage =
     ::std::variant<std::monostate, IncomingMessageFromHost, ForceUSBDisconnect,
                    ErrorMessage, AcknowledgePrevious, GetSystemInfoResponse,
@@ -285,6 +289,8 @@ using SystemMessage =
                    SetSerialNumberMessage, EnterBootloaderMessage,
                    GetDoorClosedMessage>;
 
+using UIMessage = ::std::variant<std::monostate, UpdateUIMessage>;
+
 using MotorDriverMessage =
     ::std::variant<std::monostate, SetTMCRegisterMessage, GetTMCRegisterMessage,
                    PollTMCRegisterMessage, StopPollTMCRegisterMessage,
diff --git a/stm32-modules/include/flex-stacker/flex-stacker/motor_task.hpp b/stm32-modules/include/flex-stacker/flex-stacker/motor_task.hpp
index 09c81066b..bbb4a10d7 100644
--- a/stm32-modules/include/flex-stacker/flex-stacker/motor_task.hpp
+++ b/stm32-modules/include/flex-stacker/flex-stacker/motor_task.hpp
@@ -69,7 +69,7 @@ struct Defaults {
         static constexpr float SPEED_DISCONT = 5.0;
 
         static constexpr float MM_PER_REV =
-            lms::GearBoxConfig::mm_per_rev(16.0, 16.0 / 30.0);
+            lms::GearBoxConfig::mm_per_rev(16.0, 30.0/16.0);
         static constexpr float STEPS_PER_REV = 200;
         static constexpr float MICROSTEP = 16;
     };
diff --git a/stm32-modules/include/flex-stacker/flex-stacker/tasks.hpp b/stm32-modules/include/flex-stacker/flex-stacker/tasks.hpp
index 79e85127a..fc2b91a45 100644
--- a/stm32-modules/include/flex-stacker/flex-stacker/tasks.hpp
+++ b/stm32-modules/include/flex-stacker/flex-stacker/tasks.hpp
@@ -19,10 +19,12 @@ struct Tasks {
     using HostCommsQueue = QueueImpl<messages::HostCommsMessage>;
     // Message queue for system task
     using SystemQueue = QueueImpl<messages::SystemMessage>;
+    // Message queue for UI task
+    using UIQueue = QueueImpl<messages::UIMessage>;
     // Central aggregator
     using QueueAggregator =
         queue_aggregator::QueueAggregator<MotorDriverQueue, MotorQueue,
-                                          HostCommsQueue, SystemQueue>;
+                                          HostCommsQueue, SystemQueue, UIQueue>;
 
     // Addresses
     static constexpr size_t MotorDriverAddress =
@@ -33,6 +35,8 @@ struct Tasks {
         QueueAggregator::template get_queue_idx<HostCommsQueue>();
     static constexpr size_t SystemAddress =
         QueueAggregator::template get_queue_idx<SystemQueue>();
+    static constexpr size_t UIAddress =
+        QueueAggregator::template get_queue_idx<UIQueue>();
 };
 
 };  // namespace tasks
diff --git a/stm32-modules/include/flex-stacker/flex-stacker/ui_task.hpp b/stm32-modules/include/flex-stacker/flex-stacker/ui_task.hpp
new file mode 100644
index 000000000..28d70e072
--- /dev/null
+++ b/stm32-modules/include/flex-stacker/flex-stacker/ui_task.hpp
@@ -0,0 +1,178 @@
+#pragma once
+
+#include <algorithm>
+#include <array>
+
+#include "core/is31fl_driver.hpp"
+#include "hal/message_queue.hpp"
+#include "messages.hpp"
+#include "ui_policy.hpp"
+
+namespace ui_task {
+
+template <typename Policy>
+concept UIPolicy = requires(Policy& p) {
+    // A function to set the heartbeat LED on or off
+    {p.set_heartbeat_led(true)};
+}
+&&is31fl::IS31FL_Policy<Policy>;
+
+// Structure to hold runtime info for the heartbeat LED.
+// The LED is run with a psuedo-pwm to confirm that the
+// firmware is running and tasks are being called at
+// regular intervals
+class Heartbeat {
+  private:
+    // This gives a pleasant visual effect
+    static constexpr uint32_t default_period = 25;
+    const uint32_t _period;
+    uint8_t _pwm = 0;
+    uint8_t _count = 0;
+    int8_t _direction = 1;
+
+  public:
+    explicit Heartbeat(uint32_t period = default_period) : _period(period) {}
+
+    /**
+     * @brief Increment heartbeat counter. This provides a pseudo-pwm
+     * setup where a "pwm" counter runs from 0 to a configurable period
+     * value, and the LED is turned on and off based on whether the repeating
+     * counter is below the PWM value.
+     *
+     * @return true if the LED should be set to on, false if it should
+     * be set to off.
+     */
+    auto tick() -> bool {
+        ++_count;
+        if (_count == _period) {
+            _count = 0;
+            _pwm += _direction;
+            if (_pwm == _period) {
+                _direction = -1;
+            } else if (_pwm == 0) {
+                _direction = 1;
+            }
+        }
+        return (_pwm > 2) && (_count < _pwm);
+    }
+
+    // Get the current pwm period
+    [[nodiscard]] auto pwm() const -> uint8_t { return _pwm; }
+};
+
+enum Color { W, R, G, B };
+
+// There are 3 channels per color
+using ChannelMapping = std::array<size_t, 3>;
+
+static constexpr ChannelMapping white_channels{3, 4, 5};
+static constexpr ChannelMapping red_channels{6, 9, 12};
+static constexpr ChannelMapping green_channels{7, 10, 13};
+static constexpr ChannelMapping blue_channels{8, 11, 14};
+
+static auto color_to_channels(Color color) -> const ChannelMapping& {
+    switch (color) {
+        case Color::W:
+            return white_channels;
+        case Color::R:
+            return red_channels;
+        case Color::G:
+            return green_channels;
+        case Color::B:
+            return blue_channels;
+        default:
+            return white_channels;
+    }
+}
+
+using Message = messages::UIMessage;
+
+template <template <class> class QueueImpl>
+requires MessageQueue<QueueImpl<Message>, Message>
+class UITask {
+  public:
+    using Queue = QueueImpl<Message>;
+    using Aggregator = typename tasks::Tasks<QueueImpl>::QueueAggregator;
+    using Queues = typename tasks::Tasks<QueueImpl>;
+
+    // The timer driving LED update frequency should run at this period
+    static constexpr uint32_t UPDATE_PERIOD_MS = 1;
+
+    static constexpr uint8_t LED_DRIVER_I2C_ADDRESS = 0x6C << 1;
+
+    explicit UITask(Queue& q, Aggregator* aggregator)
+        : _message_queue(q),
+          _task_registry(aggregator),
+          // NOLINTNEXTLINE(readability-redundant-member-init)
+          _heartbeat(),
+          // NOLINTNEXTLINE(readability-redundant-member-init)
+          _led_driver() {}
+    UITask(const UITask& other) = delete;
+    auto operator=(const UITask& other) -> UITask& = delete;
+    UITask(UITask&& other) noexcept = delete;
+    auto operator=(UITask&& other) noexcept -> UITask& = delete;
+    ~UITask() = default;
+
+    auto provide_aggregator(Aggregator* aggregator) {
+        _task_registry = aggregator;
+    }
+
+    // This should be called from the periodic update timer to drive the
+    // update frequency of the LED's on the system
+    auto update_callback() {
+        static_cast<void>(_message_queue.try_send(messages::UpdateUIMessage{}));
+    }
+
+    template <UIPolicy Policy>
+    auto run_once(Policy& policy) -> void {
+        auto message = Message(std::monostate());
+
+        if (!_led_driver.initialized()) {
+            _led_driver.initialize(policy);
+            // Set the status bar to all green
+            set_color_power(Color::G, 1);
+            _led_driver.set_pwm(1.0F);
+            _led_driver.send_update(policy);
+        }
+
+        _message_queue.recv(&message);
+        auto visit_helper = [this, &policy](auto& message) -> void {
+            this->visit_message(message, policy);
+        };
+        std::visit(visit_helper, message);
+    }
+
+  private:
+    template <UIPolicy Policy>
+    auto visit_message(const std::monostate& message, Policy& policy) -> void {
+        static_cast<void>(message);
+        static_cast<void>(policy);
+    }
+
+    template <UIPolicy Policy>
+    auto visit_message(const messages::UpdateUIMessage& message, Policy& policy)
+        -> void {
+        static_cast<void>(message);
+        policy.set_heartbeat_led(_heartbeat.tick());
+    }
+
+    /**
+     * @brief Set the power (separate from PWM) for a color. Each color has
+     * 3 channels, so this helper will set all of the channels.
+     */
+    auto set_color_power(Color color, float power) -> bool {
+        const auto& channels = color_to_channels(color);
+
+        return std::ranges::all_of(channels.cbegin(), channels.cend(),
+                                   [power, this](size_t c) {
+                                       return _led_driver.set_current(c, power);
+                                   });
+    }
+
+    Queue& _message_queue;
+    Aggregator* _task_registry;
+    Heartbeat _heartbeat;
+    is31fl::IS31FL<LED_DRIVER_I2C_ADDRESS> _led_driver;
+};
+
+};  // namespace ui_task