Merge pull request #36 from tidalf/add-stats

AlexxIT · web-flow · commit ec341cfaf9b8 · 2025-05-01T15:26:14.000+03:00
Add stats (1/3)
diff --git a/custom_components/jura/__init__.py b/custom_components/jura/__init__.py
@@ -9,7 +9,7 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-PLATFORMS = ["binary_sensor", "button", "number", "select", "switch"]
+PLATFORMS = ["binary_sensor", "button", "number", "select", "switch", "sensor"]
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
diff --git a/custom_components/jura/core/client.py b/custom_components/jura/core/client.py
@@ -10,10 +10,25 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-ACTIVE_TIME = 120
+ACTIVE_TIME = 30
 COMMAND_TIME = 15
 
 
+class UUIDs:
+    """BLE characteristic UUIDs."""
+
+    # https://github.com/Jutta-Proto/protocol-bt-cpp?tab=readme-ov-file#bluetooth-characteristics
+    # Heartbeat
+    ABOUT_MACHINE = "5a401531-ab2e-2548-c435-08c300000710"
+
+    # Start product
+    START_PRODUCT = "5a401525-ab2e-2548-c435-08c300000710"
+
+    # Statistics
+    STATS_COMMAND = "5a401533-ab2e-2548-c435-08c300000710"
+    STATS_DATA = "5A401534-ab2e-2548-c435-08c300000710"
+
+
 class Client:
     def __init__(self, device: BLEDevice, callback: Callable = None):
         self.device = device
@@ -25,9 +40,10 @@ def __init__(self, device: BLEDevice, callback: Callable = None):
         self.ping_future: asyncio.Future | None = None
         self.ping_task: asyncio.Task | None = None
         self.ping_time = 0
-
+        self.key = None
         self.send_data = None
         self.send_time = 0
+        self.send_uuid = None
 
     def ping(self):
         self.ping_time = time.time() + ACTIVE_TIME
@@ -43,10 +59,11 @@ def ping_cancel(self):
         if self.ping_future:
             self.ping_future.cancel()
 
-    def send(self, data: bytes):
+    def send(self, data: bytes, uuid: str = UUIDs.START_PRODUCT):
         # if send loop active - we change sending data
         self.send_time = time.time() + COMMAND_TIME
         self.send_data = data
+        self.send_uuid = uuid
 
         # refresh ping time
         self.ping()
@@ -68,15 +85,16 @@ async def _ping_loop(self):
                 while time.time() < self.ping_time:
                     # important dummy read for keep connection
                     data = await self.client.read_gatt_char(
-                        "5a401531-ab2e-2548-c435-08c300000710"
+                        UUIDs.ABOUT_MACHINE
                     )
-                    key = data[0]
+                    self.key = data[0]
+                    _LOGGER.debug(f"key: {self.key}")
 
                     if self.send_data:
                         if time.time() < self.send_time:
                             await self.client.write_gatt_char(
-                                "5a401525-ab2e-2548-c435-08c300000710",
-                                data=encrypt(self.send_data, key),
+                                self.send_uuid,
+                                data=encrypt(self.send_data, self.key),
                                 response=True,
                             )
                         self.send_data = None
@@ -104,6 +122,62 @@ async def _ping_loop(self):
 
         self.ping_task = None
 
+    async def read(self, uuid: str, decrypt: bool = False):
+        """Read data from a characteristic."""
+        if not self.client:
+            _LOGGER.warning("Cannot read: No active client connection")
+            return None
+
+        try:
+            data = await self.client.read_gatt_char(uuid)
+            if decrypt and self.key:
+                return encryption.encdec(list(data), self.key)
+            return data
+        except BleakError as e:
+            _LOGGER.error(
+                f"Error reading from characteristic {uuid}", exc_info=e)
+            raise
+        except Exception as e:
+            _LOGGER.error(
+                f"Error reading from characteristic {uuid}", exc_info=e)
+            raise
+
+    async def read_statistics_data(self, timeout: int = 20, retries: int = 30) -> bytes | None:
+        """Read statistics data from the device."""
+        _LOGGER.debug("Reading Jura statistics...")
+
+        # Send statistics request command
+        # https://github.com/Jutta-Proto/protocol-bt-cpp?tab=readme-ov-file#writing-1
+        command_bytes = [0x2A, 0x00, 0x01, 0xFF, 0xFF]
+        self.send(bytes(command_bytes), uuid=UUIDs.STATS_COMMAND)
+
+        # Wait for connection
+        if not self.client:
+            _LOGGER.debug("Waiting for client connection...")
+            for _ in range(timeout):
+                if not self.client:
+                    await asyncio.sleep(1)
+                else:
+                    break
+            if not self.client:
+                _LOGGER.debug("Failed to establish connection")
+                return None
+
+        # Wait for statistics to be ready
+        # https://github.com/Jutta-Proto/protocol-bt-cpp?tab=readme-ov-file#reading
+        for _ in range(retries):
+            status = await self.read(UUIDs.STATS_COMMAND)
+            if status and status[1] != 225:  # 225 means not ready
+                break
+            await asyncio.sleep(0.8)
+        else:
+            _LOGGER.error("Device not ready for statistics reading")
+            return None
+
+        # Read statistics data
+        # https://github.com/Jutta-Proto/protocol-bt-cpp?tab=readme-ov-file#statistics-data
+        return await self.read(UUIDs.STATS_DATA, decrypt=True)
+
 
 def encrypt(data: bytes, key: int) -> bytes:
     data = bytearray(data)
diff --git a/custom_components/jura/core/device.py b/custom_components/jura/core/device.py
@@ -2,12 +2,16 @@
 from pathlib import Path
 from typing import Callable, TypedDict
 from zipfile import ZipFile
+import logging
 
 import xmltodict
 from bleak import AdvertisementData, BLEDevice
 
 from .client import Client
 
+_LOGGER = logging.getLogger(__name__)
+
+COMMAND_TIME = 15
 SELECTS = [
     "product",  # 1
     "grinder_ratio",  # 2
@@ -54,6 +58,8 @@ def __init__(self, name: str, model: str, products: list, device: BLEDevice):
         self.values = None
         self.updates_connect: list = []
         self.updates_product: list = []
+        self.updates_statistics = []
+        self.statistics = {"total_products": None, "product_counts": {}}
 
     @property
     def mac(self) -> str:
@@ -129,7 +135,8 @@ def select_option(self, attr: str, option: str):
         if not attribute:
             return None
 
-        value = next(i["@Value"] for i in attribute["ITEM"] if i["@Name"] == option)
+        value = next(i["@Value"]
+                     for i in attribute["ITEM"] if i["@Name"] == option)
         self.set_value(attr, int(value, 16))
 
     def set_value(self, attr: str, value: int):
@@ -182,10 +189,83 @@ def command(self) -> bytes:
         # data[0] = self.key
         # data[9] = 1
         # data[16] = 6
-        # data[17] = self.key
+
+        # need to be set or the machine will go into a half broken state
+        data[17] = self.client.key
 
         return data
 
+    # Add method to register statistics updates
+    def register_statistics_update(self, handler: Callable):
+        """Register a callback for statistics updates."""
+        self.updates_statistics.append(handler)
+
+    async def read_statistics(self, force_update: bool = False):
+        """Read statistics from the machine."""
+
+        _LOGGER.debug("Reading Jura statistics...")
+
+        # Read statistics data from client
+        decrypted_data = await self.client.read_statistics_data()
+        if decrypted_data is None:
+            _LOGGER.debug(
+                "Failed to read statistics data, returning existing statistics")
+            return self.statistics
+
+        # Convert all 3-byte chunks to integers
+        product_counts_array = []
+        for i in range(0, len(decrypted_data), 3):
+            if i + 3 <= len(decrypted_data):
+                # Convert 3 bytes to an integer
+                count = int.from_bytes(decrypted_data[i:i+3], 'big')
+                if count == 0xFFFF:  # means 0 it seems
+                    count = 0
+                product_counts_array.append(count)
+
+        # get total_count from first 3 bytes if available
+        total_count = product_counts_array[0] if product_counts_array and product_counts_array[0] is not None else None
+        _LOGGER.info(
+            f"Total coffee count from data: {total_count if total_count is not None else 'undefined'}")
+
+        # remove aberrant values if any
+        if total_count == 0 or total_count > 1000000:
+            _LOGGER.info(
+                "total 0 or too high, something's wrong, returning existing statistics")
+            return self.statistics
+        # get the names associated to the products counts
+        product_counts = {}
+        for i, count in enumerate(product_counts_array):
+            if i == 0:  # Skip the total
+                continue
+
+            product = next(
+                (p for p in self.products if int(p['@Code'], 16) == i), None)
+            if product:
+                product_counts[product['@Name']] = count
+                _LOGGER.debug(
+                    f"Stat entry: Position {i} = {count} -> {product['@Name']}")
+            else:
+                _LOGGER.debug(
+                    f"No product found for code {i} with count {count}")
+
+        # Log the final counts at info log level
+        for product, count in product_counts.items():
+            _LOGGER.info(f"Product: {product}, Count: {count}")
+
+        # Save the statistics
+        self.statistics = {
+            "total_products": total_count,
+            "product_counts": product_counts
+        }
+
+        # Notify all statistics listeners
+        _LOGGER.debug(
+            f"Notifying {len(self.updates_statistics)} statistics listeners")
+        for handler in self.updates_statistics:
+            handler()
+
+        return self.statistics
+
 
 class EmptyModel(Exception):
     pass
diff --git a/custom_components/jura/sensor.py b/custom_components/jura/sensor.py
@@ -0,0 +1,127 @@
+"""Sensor platform for Jura integration."""
+import logging
+import asyncio
+from datetime import timedelta
+from typing import Any, Callable, Optional
+
+from homeassistant.components.sensor import (
+    SensorDeviceClass,
+    SensorEntity,
+    SensorStateClass,
+)
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.event import async_track_time_interval
+
+from .core import DOMAIN
+from .core.entity import JuraEntity
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+) -> None:
+    """Set up Jura sensor based on a config entry."""
+    device = hass.data[DOMAIN][entry.entry_id]
+
+    # Create the total coffees sensor
+    entities = [JuraTotalCoffeeSensor(device)]
+
+    # Create sensors for each product
+    for product in device.products:
+        product_name = product["@Name"]
+        if product.get("@Active") != "false":
+            entities.append(JuraProductCountSensor(device, product_name))
+
+    async_add_entities(entities)
+
+    # Set up automatic refresh
+    update_interval = hass.data[DOMAIN].get("update_interval", 60)
+
+    async def refresh_statistics(*_):
+        """Refresh statistics regularly."""
+        try:
+            await device.read_statistics()
+        except Exception as ex:
+            _LOGGER.error(f"Error refreshing statistics: {ex}")
+
+    # Schedule regular updates
+    entry.async_on_unload(
+        async_track_time_interval(
+            hass,
+            refresh_statistics,
+            timedelta(seconds=update_interval)
+        )
+    )
+
+    # Do an initial refresh
+    hass.async_create_task(refresh_statistics())
+
+
+class JuraStatisticsSensor(JuraEntity, SensorEntity):
+    """Base class for Jura statistics sensors."""
+
+    def __init__(self, device, attr: str):
+        """Initialize the sensor."""
+        super().__init__(device, attr)
+
+        # Register for updates on statistics
+        device.register_statistics_update(self.internal_update)
+
+    @property
+    def native_value(self) -> Any:
+        """Return the state of the sensor."""
+        return self._get_value()
+
+    def _get_value(self) -> Any:
+        """Get the value for this sensor from statistics."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def internal_update(self):
+        """Override parent method to ensure statistics are refreshed."""
+        _LOGGER.debug(f"Updating sensor {self._attr_name}")
+        if self.hass is not None:
+            self.async_write_ha_state()
+
+
+class JuraTotalCoffeeSensor(JuraStatisticsSensor):
+    """Sensor for total coffee count."""
+
+    _attr_icon = "mdi:coffee"
+    _attr_state_class = SensorStateClass.TOTAL_INCREASING
+    _attr_native_unit_of_measurement = "products"
+
+    def __init__(self, device):
+        """Initialize the sensor."""
+        super().__init__(device, "total_product")
+        self._attr_name = f"{device.name} Total Products"
+
+    def _get_value(self) -> int:
+        """Get the total coffee count."""
+        value = self.device.statistics.get("total_products", 0)
+        _LOGGER.debug(f"Total coffee value: {value}")
+        return value
+
+
+class JuraProductCountSensor(JuraStatisticsSensor):
+    """Sensor for individual product count."""
+
+    _attr_icon = "mdi:coffee-outline"
+    _attr_state_class = SensorStateClass.TOTAL_INCREASING
+    _attr_native_unit_of_measurement = "products"
+
+    def __init__(self, device, product_name: str):
+        """Initialize the sensor."""
+        self.product_name = product_name
+        attr_name = f"product_{product_name.lower().replace(' ', '_')}"
+        super().__init__(device, attr_name)
+        self._attr_name = f"{device.name} {product_name} Count"
+
+    def _get_value(self) -> int:
+        """Get the count for this specific product."""
+        value = self.device.statistics.get(
+            "product_counts", {}).get(self.product_name, None)
+        _LOGGER.debug(f"Product {self.product_name} count: {value}")
+        return value
