diff --git a/oin_ha/display/__init__.py b/oin_ha/display/__init__.py
index 5283925..fa628d3 100644
--- a/oin_ha/display/__init__.py
+++ b/oin_ha/display/__init__.py
@@ -102,7 +102,7 @@ class Light:
                         "retain": True,
                         "unique_id": f"{self.uid}_{i}",
                         "state_topic": self.state_topic,
-                        "state_value_template": "{{ value_json.state }}",
+                        "state_value_template": "{{" + f"value_json.state[{i}]" + "}}",
                     }
                     | self.options
                 ),
@@ -120,19 +120,19 @@ class Light:
 
     def on_message(self, client, userdata, message):
         data = message.payload.decode()
-        print(message.topic, data)
 
         match message.topic.rsplit("/", maxsplit=2):
             case [self.base_topic, i, "command"]:
                 self.switch = False
             case [self.base_topic, i, "rgb"]:
+                self._i = int(i)
                 self.set_color(int(i), list(map(int, data.split(","))))
             case [self.base_topic, i, "effect"]:
                 self.low_light = data == "Low Light"
             case [self.base_topic, i, "value"]:
                 self.set_value(int(i), data)
             case [self.base_topic, i, "action_color"]:
-                self.set_color(int(i), self._colors.get(data, [0, 0, 0]), False)
+                self.set_color(int(i), self._colors.get(data, [0, 0, 0]), switch=False)
             case _:
                 return
 
@@ -143,7 +143,9 @@ class Light:
                 {
                     "effect": "Low Light" if self.low_light else "Normal",
                     "rgb": self.rgb,
-                    "state": self.state,
+                    "state": [
+                        self.state if i == self._i else "OFF" for i in range(self._n)
+                    ],
                 }
             ),
             retain=True,
@@ -192,7 +194,9 @@ class Light:
         self._color[i] = value
         if switch and not self.switch:
             self.switch = True
-        if i == self._i:
+        if switch:
+            self.update_value()
+        elif i == self._i:
             self.display_value()
         self.publish_state()