From 3f5b4bedeec1d157d10a46cb1f8881861d7bd302 Mon Sep 17 00:00:00 2001
From: Nicolas Duhamel <nicolas@jombi.fr>
Date: Tue, 16 Mar 2021 09:08:08 +0100
Subject: [PATCH] Handle LWT

---
 src/citadel/devices/__init__.py | 38 ++++++++++++++++++++++++++++++---
 src/citadel/devices/__main__.py |  3 ++-
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/src/citadel/devices/__init__.py b/src/citadel/devices/__init__.py
index 11ab662..77e20dd 100644
--- a/src/citadel/devices/__init__.py
+++ b/src/citadel/devices/__init__.py
@@ -7,6 +7,7 @@ class Light:
     def __init__(self, _id: str, client: Client, **kwargs):
         self._state = None
         self._change_callback = None
+        self._available = False
 
         self._id = _id
 
@@ -18,12 +19,28 @@ class Light:
         self.CMND_ON_PAYLOAD = kwargs.get('CMND_ON_PAYLOAD', 'ON')
         self.CMND_OFF_PAYLOAD = kwargs.get('CMND_OFF_PAYLOAD', 'OFF')
 
+        self.LWT_TOPIC = kwargs.get('LWT_TOPIC', 'tele/tasmota/light/{id}/LWT')
+
         self._client = client
         self._client.subscribe_callback(self.STATE_TOPIC.format(id=self.id), self.onStateChange)
+        self._client.subscribe_callback(self.LWT_TOPIC.format(id=self.id), self.onLWTChange)
 
     def setChangeCallback(self, callback: typing.Callable):
         self._change_callback = callback
 
+    def onLWTChange(self, client, userdata, msg):
+        state = msg.payload.decode()
+        if state == "Online":
+            self._available = True
+        elif state == "Offline":
+            self._available = False
+            self._state = False
+        else:
+            raise Exception("Unexpected value %s for LWT topic" % state)
+
+        if self._change_callback:
+            self._change_callback(self)
+
     def onStateChange(self, client, userdata, msg):
         state = msg.payload.decode()
 
@@ -47,6 +64,10 @@ class Light:
     def state(self):
         return self._state
 
+    @property
+    def available(self):
+        return self._available
+
     @property
     def id(self):
         return self._id
@@ -58,7 +79,7 @@ class LightsGroup:
         """ Create a virtual mqtt device that handle multiple light """
         self._id = _id
         self._lights = lights
-        
+
         for l in self._lights:
             l.setChangeCallback(self.onLightStateChange)
 
@@ -66,6 +87,12 @@ class LightsGroup:
         self._client.subscribe_callback('cmnd/devices/light/%s/POWER' % self._id, self.cmndPower)
 
         self._previous_state = all([l.state for l in self._lights])
+        self._previous_available = all([l.available for l in self._lights])
+
+        self._client.add_on_connect_callback(self.on_connect)
+
+    def on_connect(self, *args):
+        self._client.publish('stat/devices/light/%s/AVAILABLE' % self.id, 'TRUE' if self._previous_available else 'FALSE')
 
     def cmndPower(self, client, userdata, msg):
         value = msg.payload.decode()
@@ -79,11 +106,16 @@ class LightsGroup:
             raise Exception("Unexpected value %s for stat topic" % value)
 
     def onLightStateChange(self, light):
+        available = all([l.available for l in self._lights])
+        if available != self._previous_available:
+            self._previous_available = available
+            self._client.publish('stat/devices/light/%s/AVAILABLE' % self.id, 'TRUE' if available else 'FALSE')
+
         state = all([l.state for l in self._lights])
         if state != self._previous_state:
             self._client.publish('stat/devices/light/%s/POWER' % self.id, 'ON' if state else 'OFF')
             self._previous_state = state
-    
+
     def on(self):
         for l in self._lights:
             l.on()
@@ -95,7 +127,7 @@ class LightsGroup:
     @property
     def state(self):
         return self._previous_state
-    
+
     @property
     def id(self):
         return self._id
diff --git a/src/citadel/devices/__main__.py b/src/citadel/devices/__main__.py
index c5d20a7..2e2059c 100644
--- a/src/citadel/devices/__main__.py
+++ b/src/citadel/devices/__main__.py
@@ -36,7 +36,6 @@ def main(\
                     STATE_TOPIC='stat/tasmota/light/{id}/POWER2',
                     CMND_TOPIC='cmnd/tasmota/light/{id}/POWER2')])
 
-    client.loop_forever()
 
     if is_systemd:
         logger = logging.getLogger('')
@@ -55,5 +54,7 @@ def main(\
         # add ch to logger
         logger.addHandler(ch)
 
+    client.loop_forever()
+
 if __name__ == "__main__":
     typer.run(main)