diff --git a/doc/example.yaml b/doc/example.yaml
index 731df7c..7a55709 100644
--- a/doc/example.yaml
+++ b/doc/example.yaml
@@ -67,7 +67,11 @@ light:
       then:
         - output.set_level:
             id: ${id_front_panel_output}
-            level: !lambda if (x < 0.012f) return 0; else return x;
+            level: !lambda |-
+              if (id(${id_light_mode}).state == "night")
+                return 0;
+              else
+                return x;
     # You can use any effects that you like. These are just examples.
     effects:
       - random:
diff --git a/light/__init__.py b/light/__init__.py
index 9bf1ab7..1f676f2 100644
--- a/light/__init__.py
+++ b/light/__init__.py
@@ -17,7 +17,6 @@ CONF_ON_BRIGHTNESS = "on_brightness"
 YeelightBS2LightState = bs2_ns.class_("YeelightBS2LightState", light.LightState)
 YeelightBS2LightOutput = bs2_ns.class_("YeelightBS2LightOutput", light.LightOutput)
 BrightnessTrigger = bs2_ns.class_("BrightnessTrigger", automation.Trigger.template())
-LightModeTrigger = bs2_ns.class_("LightModeTrigger", automation.Trigger.template())
 
 CONFIG_SCHEMA = light.RGB_LIGHT_SCHEMA.extend(
     {
diff --git a/light/automation.h b/light/automation.h
index aca3c02..594e448 100644
--- a/light/automation.h
+++ b/light/automation.h
@@ -12,38 +12,20 @@ namespace bs2 {
 class BrightnessTrigger : public Trigger<float> {
 public:
     explicit BrightnessTrigger(YeelightBS2LightOutput *parent) {
-        parent->add_on_state_callback(
-            [this](light::LightColorValues values, std::string light_mode) {
-                auto new_brightness = values.get_brightness();
-                if (values.get_state() == 0) {
-                    new_brightness = 0.0f;
-                }
-                new_brightness = roundf(new_brightness * 100.0f) / 100.0f;
-                if (last_brightness_ != new_brightness) {
-                    trigger(new_brightness);
-                    last_brightness_ = new_brightness;
-                }
+        parent->add_on_state_callback([this](light::LightColorValues values) {
+            auto new_brightness = values.get_brightness();
+            if (values.get_state() == 0) {
+                new_brightness = 0.0f;
             }
-        );
-    }
-protected:
-    float last_brightness_ = -1.0f;
-};
-
-class LightModeTrigger : public Trigger<std::string> {
-public:
-    explicit LightModeTrigger(YeelightBS2LightOutput *parent) {
-        parent->add_on_state_callback(
-            [this](light::LightColorValues values, std::string light_mode) {
-                if (last_light_mode_ != light_mode) {
-                    trigger(light_mode);
-                    last_light_mode_ = light_mode;
-                }
+            new_brightness = roundf(new_brightness * 100.0f) / 100.0f;
+            if (last_brightness_ != new_brightness) {
+                trigger(new_brightness);
+                last_brightness_ = new_brightness;
             }
-        );
+        });
     }
 protected:
-    std::string last_light_mode_ = LIGHT_MODE_UNKNOWN;
+    float last_brightness_ = -1.0f;
 };
 
 } // namespace yeelight_bs2
diff --git a/light/light_output.h b/light/light_output.h
index 06617e4..81f04c3 100644
--- a/light/light_output.h
+++ b/light/light_output.h
@@ -39,7 +39,11 @@ public:
         return traits;
     }
 
-    void add_on_state_callback(std::function<void(light::LightColorValues, std::string)> &&callback) {
+    void add_on_light_mode_callback(std::function<void(std::string)> &&callback) {
+          light_mode_callback_.add(std::move(callback));
+    }
+
+    void add_on_state_callback(std::function<void(light::LightColorValues)> &&callback) {
           state_callback_.add(std::move(callback));
     }
 
@@ -57,17 +61,15 @@ public:
         GPIOOutputs *delegate;
         if (transition_handler_->set_light_color_values(values)) {
             delegate = transition_handler_;
-            state_callback_.call(
-                transition_handler_->get_end_values(),
-                delegate->light_mode);
+            light_mode_callback_.call(delegate->light_mode);
+            state_callback_.call(transition_handler_->get_end_values());
         } else {
             instant_handler_->set_light_color_values(values);
             delegate = instant_handler_;
-            state_callback_.call(values, delegate->light_mode);
+            light_mode_callback_.call(delegate->light_mode);
+            state_callback_.call(values);
         }
 
-        light_mode_ = delegate->light_mode;
-
         // Note: one might think that it is more logical to turn on the LED
         // circuitry master switch after setting the individual channels,
         // but this is the order that was used by the original firmware. I
@@ -88,16 +90,12 @@ public:
             light_->turn_off();
     }
 
-    std::string get_light_mode() {
-        return light_mode_;
-    }
-
 protected:
     LightHAL *light_;
     ColorTransitionHandler *transition_handler_;
     ColorInstantHandler *instant_handler_ = new ColorInstantHandler();
-    CallbackManager<void(light::LightColorValues, std::string)> state_callback_{};
-    std::string light_mode_;
+    CallbackManager<void(std::string)> light_mode_callback_{};
+    CallbackManager<void(light::LightColorValues)> state_callback_{};
 
     friend class YeelightBS2LightState;
 
diff --git a/text_sensor/light_mode_text_sensor.h b/text_sensor/light_mode_text_sensor.h
index 91d743e..329f4f5 100644
--- a/text_sensor/light_mode_text_sensor.h
+++ b/text_sensor/light_mode_text_sensor.h
@@ -1,8 +1,5 @@
 #pragma once 
 
-#include <cmath>
-#include "../common.h"
-#include "../front_panel_hal.h"
 #include "esphome/components/text_sensor/text_sensor.h"
 
 namespace esphome {
@@ -20,14 +17,12 @@ public:
     void set_parent(YeelightBS2LightOutput *light) { light_ = light; }
 
     void setup() {
-        light_->add_on_state_callback(
-            [this](light::LightColorValues values, std::string light_mode) {
-                if (last_light_mode_ != light_mode) {
-                    publish_state(light_mode);
-                    last_light_mode_ = light_mode;
-                }
+        light_->add_on_light_mode_callback([this](std::string light_mode) {
+            if (last_light_mode_ != light_mode) {
+                publish_state(light_mode);
+                last_light_mode_ = light_mode;
             }
-        );
+        });
     }
 
 protected: