From 558588ee8ae3e0b8f95d9a8f53b23dd92031a047 Mon Sep 17 00:00:00 2001
From: ChuckMash <86080247+ChuckMash@users.noreply.github.com>
Date: Tue, 6 Feb 2024 17:41:40 -0800
Subject: [PATCH] WLED Sync fix and BK72XX support (#6190)

Co-authored-by: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
---
 esphome/components/wled/__init__.py           |  7 ++-
 esphome/components/wled/wled_light_effect.cpp | 43 +++++++++++++++++--
 esphome/components/wled/wled_light_effect.h   |  4 ++
 3 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/esphome/components/wled/__init__.py b/esphome/components/wled/__init__.py
index 2795529203..396d5891d8 100644
--- a/esphome/components/wled/__init__.py
+++ b/esphome/components/wled/__init__.py
@@ -8,6 +8,8 @@ wled_ns = cg.esphome_ns.namespace("wled")
 WLEDLightEffect = wled_ns.class_("WLEDLightEffect", AddressableLightEffect)
 
 CONFIG_SCHEMA = cv.All(cv.Schema({}), cv.only_with_arduino)
+CONF_SYNC_GROUP_MASK = "sync_group_mask"
+CONF_BLANK_ON_START = "blank_on_start"
 
 
 @register_addressable_effect(
@@ -16,10 +18,13 @@ CONFIG_SCHEMA = cv.All(cv.Schema({}), cv.only_with_arduino)
     "WLED",
     {
         cv.Optional(CONF_PORT, default=21324): cv.port,
+        cv.Optional(CONF_SYNC_GROUP_MASK, default=0): cv.int_range(min=0, max=255),
+        cv.Optional(CONF_BLANK_ON_START, default=True): cv.boolean,
     },
 )
 async def wled_light_effect_to_code(config, effect_id):
     effect = cg.new_Pvariable(effect_id, config[CONF_NAME])
     cg.add(effect.set_port(config[CONF_PORT]))
-
+    cg.add(effect.set_sync_group_mask(config[CONF_SYNC_GROUP_MASK]))
+    cg.add(effect.set_blank_on_start(config[CONF_BLANK_ON_START]))
     return effect
diff --git a/esphome/components/wled/wled_light_effect.cpp b/esphome/components/wled/wled_light_effect.cpp
index 8c68bca6e3..7a82aaeb46 100644
--- a/esphome/components/wled/wled_light_effect.cpp
+++ b/esphome/components/wled/wled_light_effect.cpp
@@ -13,6 +13,10 @@
 #include <WiFiUdp.h>
 #endif
 
+#ifdef USE_BK72XX
+#include <WiFiUdp.h>
+#endif
+
 namespace esphome {
 namespace wled {
 
@@ -29,7 +33,11 @@ WLEDLightEffect::WLEDLightEffect(const std::string &name) : AddressableLightEffe
 void WLEDLightEffect::start() {
   AddressableLightEffect::start();
 
-  blank_at_ = 0;
+  if (this->blank_on_start_) {
+    this->blank_at_ = 0;
+  } else {
+    this->blank_at_ = UINT32_MAX;
+  }
 }
 
 void WLEDLightEffect::stop() {
@@ -101,8 +109,11 @@ bool WLEDLightEffect::parse_frame_(light::AddressableLight &it, const uint8_t *p
         if (!parse_drgb_frame_(it, payload, size))
           return false;
       } else {
-        if (!parse_notifier_frame_(it, payload, size))
+        if (!parse_notifier_frame_(it, payload, size)) {
           return false;
+        } else {
+          timeout = UINT8_MAX;
+        }
       }
       break;
 
@@ -143,8 +154,32 @@ bool WLEDLightEffect::parse_frame_(light::AddressableLight &it, const uint8_t *p
 }
 
 bool WLEDLightEffect::parse_notifier_frame_(light::AddressableLight &it, const uint8_t *payload, uint16_t size) {
-  // Packet needs to be empty
-  return size == 0;
+  // Receive at least RGBW and Brightness for all LEDs from WLED Sync Notification
+  // https://kno.wled.ge/interfaces/udp-notifier/
+  // https://github.com/Aircoookie/WLED/blob/main/wled00/udp.cpp
+
+  if (size < 34) {
+    return false;
+  }
+
+  uint8_t payload_sync_group_mask = payload[34];
+
+  if ((payload_sync_group_mask & this->sync_group_mask_) != this->sync_group_mask_) {
+    ESP_LOGD(TAG, "sync group mask does not match");
+    return false;
+  }
+
+  uint8_t bri = payload[0];
+  uint8_t r = esp_scale8(payload[1], bri);
+  uint8_t g = esp_scale8(payload[2], bri);
+  uint8_t b = esp_scale8(payload[3], bri);
+  uint8_t w = esp_scale8(payload[8], bri);
+
+  for (auto &&led : it) {
+    led.set(Color(r, g, b, w));
+  }
+
+  return true;
 }
 
 bool WLEDLightEffect::parse_warls_frame_(light::AddressableLight &it, const uint8_t *payload, uint16_t size) {
diff --git a/esphome/components/wled/wled_light_effect.h b/esphome/components/wled/wled_light_effect.h
index 8f239276d7..a591e1fd1a 100644
--- a/esphome/components/wled/wled_light_effect.h
+++ b/esphome/components/wled/wled_light_effect.h
@@ -21,6 +21,8 @@ class WLEDLightEffect : public light::AddressableLightEffect {
   void stop() override;
   void apply(light::AddressableLight &it, const Color &current_color) override;
   void set_port(uint16_t port) { this->port_ = port; }
+  void set_sync_group_mask(uint8_t mask) { this->sync_group_mask_ = mask; }
+  void set_blank_on_start(bool blank) { this->blank_on_start_ = blank; }
 
  protected:
   void blank_all_leds_(light::AddressableLight &it);
@@ -35,6 +37,8 @@ class WLEDLightEffect : public light::AddressableLightEffect {
   std::unique_ptr<UDP> udp_;
   uint32_t blank_at_{0};
   uint32_t dropped_{0};
+  uint8_t sync_group_mask_{0};
+  bool blank_on_start_{true};
 };
 
 }  // namespace wled