Add files via upload

custom_components/saas/sensor.py

@@ -88,7 +88,7 @@ class SAASSensor(RestoreEntity):
         self._hass.states.async_set(self.entity_id, self._state)
         _LOGGER.info(f"{datetime.now().strftime('%H:%M:%S:%f')} (Line {inspect.currentframe().f_lineno}): Saved state: {self._state} for sensor {self.name}")
 
-class SAASAlarmEventSensor(RestoreEntity):
+class SAASAlarmEventSensor(RestoreEntity):
     """Representation of a SAAS - Sleep As Android Stats sensor for Alarm Events."""
 
     def __init__(self, hass, name, mapping, entry_id):
@@ -215,67 +215,123 @@ class SAASAlarmEventSensor(RestoreEntity):
         """Set the state to 'None' after a timeout."""
         await asyncio.sleep(15)
         self._state = "None"
-        self._last_event = "None"
+        self._last_event = "None"
-        _LOGGER.debug(f"{datetime.now().strftime('%H:%M:%S:%f')} (Line {inspect.currentframe().f_lineno}): Set state to 'None' due to timeout for sensor {self.name}")
+        _LOGGER.debug(f"{datetime.now().strftime('%H:%M:%S:%f')} (Line {inspect.currentframe().f_lineno}): Set state to 'None' due to timeout for sensor {self.name}")
-        self.async_schedule_update_ha_state()
+        self.async_schedule_update_ha_state()
-
+
-class SAASNextAlarmSensor(RestoreEntity):
+class SAASNextAlarmSensor(RestoreEntity):
-    """Sensor that exposes the next scheduled alarm time and label."""
+    """Sensor that exposes the next scheduled alarm time and label."""
-
+
-    def __init__(self, hass, name, entry_id):
+    def __init__(self, hass, name, entry_id):
-        self._hass = hass
+        self._hass = hass
-        self._name = name
+        self._name = name
-        self.entry_id = entry_id
+        self.entry_id = entry_id
-        self._state = None
+        self._state = None
-        self._label = None
+        self._label = None
-
+        self._alarms = []
-    @property
+
-    def unique_id(self):
+    @property
-        return f"saas_next_alarm_{self._name}"
+    def unique_id(self):
-
+        return f"saas_next_alarm_{self._name}"
-    @property
+
-    def name(self):
+    @property
-        return f"SAAS {self._name} Next Alarm"
+    def name(self):
-
+        return f"SAAS {self._name} Next Alarm"
-    @property
+
-    def state(self):
+    @property
-        return self._state
+    def state(self):
-
+        return self._state
-    @property
+
-    def extra_state_attributes(self):
+    @property
-        return {"Label": self._label} if self._label else {}
+    def device_info(self):
-
+        """Return information about the device."""
-    async def async_added_to_hass(self):
+        return {
-        await super().async_added_to_hass()
+            "identifiers": {(DOMAIN, self._name)},
-
+            "name": self._name,
-        state = await self.async_get_last_state()
+            "manufacturer": INTEGRATION_NAME,
-        if state:
+            "model": MODEL,
-            self._state = state.state
+        }
-            self._label = state.attributes.get("Label")
+
-
+    @property
-        async def message_received(msg):
+    def extra_state_attributes(self):
-            msg_json = json.loads(msg.payload)
+        attrs = {}
-            event = msg_json.get("event")
+        if self._label:
-            if event != "alarm_rescheduled":
+            attrs["Label"] = self._label
-                return
+        if self._alarms:
-
+            attrs["Alarms"] = self._alarms
-            value1 = msg_json.get("value1")
+        return attrs
-            value2 = msg_json.get("value2")
+
-            if value1:
+    async def async_added_to_hass(self):
-                timestamp = int(value1) / 1000.0
+        await super().async_added_to_hass()
-                dt = dt_util.as_local(datetime.fromtimestamp(timestamp))
+
-                self._state = dt.strftime("%Y-%m-%d %H:%M")
+        state = await self.async_get_last_state()
-            else:
+        if state:
-                self._state = None
+            self._state = state.state if state.state != "None" else None
-
+            self._label = state.attributes.get("Label")
-            self._label = value2
+            saved = state.attributes.get("Alarms")
-            self.async_schedule_update_ha_state()
+            if isinstance(saved, list):
-
+                self._alarms = saved
-        await async_subscribe(
+
-            self._hass,
+        async def message_received(msg):
-            self._hass.data[DOMAIN][self.entry_id][CONF_TOPIC],
+            msg_json = json.loads(msg.payload)
-            message_received,
+            event = msg_json.get("event")
-        )
+            value1 = msg_json.get("value1")
+            value2 = msg_json.get("value2")
+
+            if event == "alarm_rescheduled":
+                if value1:
+                    ts = int(value1) / 1000.0
+                    dt = dt_util.as_local(datetime.fromtimestamp(ts))
+                    epoch = dt.timestamp()
+                    found = False
+                    for alarm in self._alarms:
+                        if alarm["timestamp"] == epoch:
+                            alarm["label"] = value2
+                            found = True
+                            break
+                    if not found:
+                        self._alarms.append({"timestamp": epoch, "label": value2})
+                        self._alarms.sort(key=lambda x: x["timestamp"])
+                        self._alarms = self._alarms[:10]
+
+            elif event in ("alarm_alert_dismiss", "alarm_skip_next"):
+                if value1:
+                    ts = int(value1) / 1000.0
+                    dt = dt_util.as_local(datetime.fromtimestamp(ts))
+                    epoch = dt.timestamp()
+                    self._alarms = [a for a in self._alarms if a["timestamp"] != epoch]
+                elif self._alarms:
+                    self._alarms.pop(0)
+            else:
+                return
+
+            if self._alarms:
+                next_alarm = self._alarms[0]
+                dt = dt_util.as_local(datetime.fromtimestamp(next_alarm["timestamp"]))
+                self._state = dt.strftime("%Y-%m-%d %H:%M")
+                self._label = next_alarm.get("label")
+            else:
+                self._state = None
+                self._label = None
+
+            self.async_schedule_update_ha_state()
+
+        await async_subscribe(
+            self._hass,
+            self._hass.data[DOMAIN][self.entry_id][CONF_TOPIC],
+            message_received,
+        )
+
+    async def async_will_remove_from_hass(self):
+        """Run when entity will be removed from hass."""
+        self.hass.states.async_set(
+            self.entity_id,
+            self._state,
+            {"Label": self._label, "Alarms": self._alarms},
+        )
+        _LOGGER.info(
+            f"{datetime.now().strftime('%H:%M:%S:%f')} (Line {inspect.currentframe().f_lineno}): Saved state: {self._state} for sensor {self.name}"
+        )
 
 class SAASSoundSensor(RestoreEntity):
     """Representation of a SAAS - Sleep As Android Stats sensor for Sound Events."""
@@ -875,11 +931,11 @@ async def async_setup_entry(hass, entry, async_add_entities):
     entry_id = entry.entry_id
     hass.data[DOMAIN][entry.entry_id] = entry.data
 
-    entities = [
+    entities = [
-        SAASSensor(hass, name, STATE_MAPPING, entry_id),
+        SAASSensor(hass, name, STATE_MAPPING, entry_id),
-        SAASAlarmEventSensor(hass, name, ALARM_EVENT_MAPPING, entry_id),
+        SAASAlarmEventSensor(hass, name, ALARM_EVENT_MAPPING, entry_id),
-        SAASNextAlarmSensor(hass, name, entry_id),
+        SAASNextAlarmSensor(hass, name, entry_id),
-        SAASSoundSensor(hass, name, SOUND_MAPPING, entry_id),
+        SAASSoundSensor(hass, name, SOUND_MAPPING, entry_id),
         SAASSleepTrackingSensor(hass, name, SLEEP_TRACKING_MAPPING, entry_id), 
         SAASDisturbanceSensor(hass, name, DISTURBANCE_MAPPING, entry_id),
         SAASLullabySensor(hass, name, LULLABY_MAPPING, entry_id),
@@ -891,4 +947,4 @@ async def async_setup_entry(hass, entry, async_add_entities):
         if hasattr(entity, "async_setup"):
             await entity.async_setup()
 
-    async_add_entities(entities)
+    async_add_entities(entities)