zha quirk for Aqara W500 lumi.airrtc.aeu001 with preset temperature support

aqara_w500.py

"""Quirk for Aqara W500 lumi.airrtc.aeu001 with preset temperature support."""

from zigpy import types
from zigpy.quirks import CustomCluster
from zigpy.quirks.v2 import (
    NumberDeviceClass,
    QuirkBuilder,
    SensorDeviceClass,
    SensorStateClass,
)
from zigpy.quirks.v2.homeassistant import UnitOfTemperature, EntityType
from zigpy.zcl.clusters.general import OnOff
from zigpy.zcl.clusters.hvac import Thermostat
from zigpy.zcl.clusters.measurement import TemperatureMeasurement
from zigpy.zcl.clusters.smartenergy import Metering
from zigpy.zcl.clusters.homeautomation import ElectricalMeasurement
from zigpy.zcl.foundation import BaseAttributeDefs, ZCLAttributeDef
import zigpy.types as t
import logging

_LOGGER = logging.getLogger(__name__)

XIAOMI_CLUSTER_ID = 0xFCC0
XIAOMI_MANF_ID = 0x115F

# Preset mode IDs (as used in the packed data)
PRESET_HOME = 1
PRESET_AWAY = 2
PRESET_SLEEP = 3
PRESET_VACATION = 5
PRESET_EVENING = 6


class AqaraW500PresetMode(types.enum8):
    """Preset modes."""
    Home = 1
    Away = 2
    Sleep = 3
    Vacation = 5
    Evening = 6
    Manual = 255


class AqaraW500SensorType(types.enum8):
    """Sensor Type."""
    Internal = 0
    External = 1
    NTC = 2


class AqaraW500NtcType(types.enum32):
    """NTC Type."""
    _10_kOhm = 10
    _50_kOhm = 50
    _100_kOhm = 100
    Unknown = 10000


class AqaraW500ThermostatCluster(CustomCluster, Thermostat):
    """Aqara W500 Custom Thermostat Cluster."""
    _CONSTANT_ATTRIBUTES = {
        Thermostat.AttributeDefs.ctrl_sequence_of_oper.id: Thermostat.ControlSequenceOfOperation.Heating_Only
    }


class AqaraW500ElectricalMeasurementCluster(CustomCluster, ElectricalMeasurement):
    """Aqara W500 Custom ElectricalMeasurement Cluster."""
    _CONSTANT_ATTRIBUTES = {
        ElectricalMeasurement.AttributeDefs.power_multiplier.id: 1,
        ElectricalMeasurement.AttributeDefs.power_divisor.id: 1,
        ElectricalMeasurement.AttributeDefs.ac_power_multiplier.id: 1,
        ElectricalMeasurement.AttributeDefs.ac_power_divisor.id: 1,
    }


class AqaraW500ManufacturerCluster(CustomCluster):
    """Aqara W500 Manufacturer Specific Cluster with preset temperature support."""

    cluster_id = XIAOMI_CLUSTER_ID

    # Attribute IDs
    WINDOW_DETECTION_ATTR = 0x0273
    CHILD_LOCK_ATTR = 0x0277
    SENSOR_TYPE_ATTR = 0x0280
    HYSTERESIS_ATTR = 0x030C
    STATE_ATTR = 0x0310
    PRESET_MODE_ATTR = 0x0311
    NTC_TYPE_ATTR = 0x0315
    PRESET_TEMPS_PACKED_ATTR = 0x0317

    # Virtual attributes for individual preset temperatures
    HOME_PRESET_TEMP_ATTR = 0x1001
    AWAY_PRESET_TEMP_ATTR = 0x1002
    SLEEP_PRESET_TEMP_ATTR = 0x1003
    VACATION_PRESET_TEMP_ATTR = 0x1005
    EVENING_PRESET_TEMP_ATTR = 0x1006

    class AttributeDefs(BaseAttributeDefs):
        """Attribute definitions."""
        window_detection = ZCLAttributeDef(
            id=0x0273, type=t.uint8_t, is_manufacturer_specific=True
        )
        child_lock = ZCLAttributeDef(
            id=0x0277, type=t.uint8_t, is_manufacturer_specific=True
        )
        sensor_type = ZCLAttributeDef(
            id=0x0280, type=t.uint8_t, is_manufacturer_specific=True
        )
        hysteresis = ZCLAttributeDef(
            id=0x030C, type=t.uint8_t, is_manufacturer_specific=True
        )
        state = ZCLAttributeDef(
            id=0x0310, type=t.uint8_t, is_manufacturer_specific=True
        )
        preset_mode = ZCLAttributeDef(
            id=0x0311, type=t.uint8_t, is_manufacturer_specific=True
        )
        ntc_type = ZCLAttributeDef(
            id=0x0315, type=t.uint32_t, is_manufacturer_specific=True
        )
        preset_temps_packed = ZCLAttributeDef(
            id=0x0317, type=t.LVBytes, is_manufacturer_specific=True
        )
        home_preset_temperature = ZCLAttributeDef(
            id=0x1001, type=t.uint16_t, is_manufacturer_specific=True
        )
        away_preset_temperature = ZCLAttributeDef(
            id=0x1002, type=t.uint16_t, is_manufacturer_specific=True
        )
        sleep_preset_temperature = ZCLAttributeDef(
            id=0x1003, type=t.uint16_t, is_manufacturer_specific=True
        )
        vacation_preset_temperature = ZCLAttributeDef(
            id=0x1005, type=t.uint16_t, is_manufacturer_specific=True
        )
        evening_preset_temperature = ZCLAttributeDef(
            id=0x1006, type=t.uint16_t, is_manufacturer_specific=True
        )

    # Map virtual attribute IDs to preset mode IDs
    VIRTUAL_ATTR_TO_MODE = {
        0x1001: PRESET_HOME,
        0x1002: PRESET_AWAY,
        0x1003: PRESET_SLEEP,
        0x1005: PRESET_VACATION,
        0x1006: PRESET_EVENING,
    }

    MODE_TO_VIRTUAL_ATTR = {v: k for k, v in VIRTUAL_ATTR_TO_MODE.items()}

    def _parse_preset_temps(self, data: bytes) -> dict:
        """Parse packed preset temperatures."""
        temps = {}
        if len(data) < 6:
            return temps
        num_presets = data[0]  # First byte is the count (0x05 = 5 presets)
        pos = 1
        for _ in range(num_presets):  # Only parse the expected number of records
            if pos + 5 > len(data):
                break
            mode_id = data[pos]
            temp = data[pos + 3] | (data[pos + 4] << 8)
            if mode_id in (PRESET_HOME, PRESET_AWAY, PRESET_SLEEP, PRESET_VACATION, PRESET_EVENING):
                temps[mode_id] = temp
            pos += 5
        return temps

    def _build_preset_temps(self, temps: dict) -> bytes:
        """Build packed preset temperatures byte array."""
        data = bytearray([0x05])
        for mode_id in [PRESET_HOME, PRESET_AWAY, PRESET_SLEEP, PRESET_VACATION, PRESET_EVENING]:
            if mode_id in temps:
                temp = temps[mode_id]
                data.extend([mode_id, 0x00, 0x00, temp & 0xFF, (temp >> 8) & 0xFF])
        data.extend([0x06, 0x00, 0x00, 0x00, 0x00])
        return bytes(data)

    def _update_attribute(self, attrid, value):
        """Handle attribute updates - parse packed data into virtual attributes."""
        super()._update_attribute(attrid, value)

        # When packed preset temps are updated, parse and update virtual attributes
        if attrid == self.PRESET_TEMPS_PACKED_ATTR and isinstance(value, (bytes, bytearray)):
            temps = self._parse_preset_temps(value)
            for mode_id, temp in temps.items():
                virtual_attr = self.MODE_TO_VIRTUAL_ATTR.get(mode_id)
                if virtual_attr:
                    super()._update_attribute(virtual_attr, temp)

    async def read_attributes(self, attributes, allow_cache=False, only_cache=False, manufacturer=None):
        """Override to handle virtual preset temperature attributes."""
        virtual_attrs = set()
        real_attrs = []

        for attr in attributes:
            if isinstance(attr, int):
                attr_id = attr
            elif hasattr(attr, 'id'):
                attr_id = attr.id
            elif isinstance(attr, str):
                attr_def = getattr(self.AttributeDefs, attr, None)
                attr_id = attr_def.id if attr_def else None
            else:
                attr_id = None

            if attr_id in self.VIRTUAL_ATTR_TO_MODE:
                virtual_attrs.add(attr_id)
            else:
                real_attrs.append(attr)

        result = {}
        failures = {}

        # Read real attributes
        if real_attrs:
            success, fail = await super().read_attributes(
                real_attrs, allow_cache=allow_cache, only_cache=only_cache, manufacturer=manufacturer
            )
            result.update(success)
            failures.update(fail)

        # Handle virtual attributes - read packed data and parse
        if virtual_attrs:
            try:
                success, fail = await super().read_attributes(
                    [self.PRESET_TEMPS_PACKED_ATTR], allow_cache=False, only_cache=False, manufacturer=XIAOMI_MANF_ID
                )
                if self.PRESET_TEMPS_PACKED_ATTR in success:
                    packed_data = success[self.PRESET_TEMPS_PACKED_ATTR]
                    if isinstance(packed_data, (bytes, bytearray)):
                        temps = self._parse_preset_temps(packed_data)
                        for attr_id in virtual_attrs:
                            mode_id = self.VIRTUAL_ATTR_TO_MODE.get(attr_id)
                            if mode_id and mode_id in temps:
                                result[attr_id] = temps[mode_id]
                            else:
                                failures[attr_id] = Exception("Preset not found")
                    else:
                        for attr_id in virtual_attrs:
                            failures[attr_id] = Exception("Invalid data type")
                else:
                    for attr_id in virtual_attrs:
                        failures[attr_id] = fail.get(self.PRESET_TEMPS_PACKED_ATTR, Exception("Read failed"))
            except Exception as e:
                _LOGGER.error("Failed to read preset temperatures: %s", e)
                for attr_id in virtual_attrs:
                    failures[attr_id] = e

        return result, failures

    async def write_attributes(self, attributes, manufacturer=None):
        """Override to handle virtual preset temperature attributes."""
        virtual_writes = {}
        real_attrs = {}

        for attr, value in attributes.items():
            if isinstance(attr, int):
                attr_id = attr
            elif hasattr(attr, 'id'):
                attr_id = attr.id
            elif isinstance(attr, str):
                attr_def = getattr(self.AttributeDefs, attr, None)
                attr_id = attr_def.id if attr_def else None
            else:
                attr_id = None

            if attr_id in self.VIRTUAL_ATTR_TO_MODE:
                virtual_writes[attr_id] = value
            else:
                real_attrs[attr] = value

        # Write real attributes
        if real_attrs:
            result = await super().write_attributes(real_attrs, manufacturer=manufacturer)
            if not virtual_writes:
                return result

        # Handle virtual attribute writes
        if virtual_writes:
            # Read current packed data
            success, _ = await super().read_attributes(
                [self.PRESET_TEMPS_PACKED_ATTR], allow_cache=False, only_cache=False, manufacturer=XIAOMI_MANF_ID
            )

            if self.PRESET_TEMPS_PACKED_ATTR in success:
                packed_data = success[self.PRESET_TEMPS_PACKED_ATTR]
                if isinstance(packed_data, (bytes, bytearray)):
                    temps = self._parse_preset_temps(packed_data)

                    # Update with new values
                    for attr_id, value in virtual_writes.items():
                        mode_id = self.VIRTUAL_ATTR_TO_MODE.get(attr_id)
                        if mode_id:
                            raw_value = int(float(value))
                            # Detect if value is in degrees (small number) vs centidegrees (large number)
                            if raw_value < 100:
                                raw_value = raw_value * 100
                            temps[mode_id] = raw_value

                    # Build and write new packed data
                    new_packed = self._build_preset_temps(temps)
                    return await super().write_attributes(
                        {self.PRESET_TEMPS_PACKED_ATTR: new_packed}, manufacturer=XIAOMI_MANF_ID
                    )

        # Return empty success if nothing to write
        from zigpy.zcl import foundation
        return [foundation.WriteAttributesStatusRecord(status=foundation.Status.SUCCESS)]


(
    QuirkBuilder("Aqara", "lumi.airrtc.aeu001")
    .prevent_default_entity_creation(
        endpoint_id=1,
        cluster_id=TemperatureMeasurement.cluster_id,
    )
    .prevent_default_entity_creation(
        endpoint_id=1,
        cluster_id=OnOff.cluster_id,
    )
    .prevent_default_entity_creation(
        endpoint_id=1,
        cluster_id=Metering.cluster_id,
    )
    .replaces(AqaraW500ElectricalMeasurementCluster)
    .prevent_default_entity_creation(
        endpoint_id=1,
        cluster_id=AqaraW500ElectricalMeasurementCluster.cluster_id,
        function=lambda entity: entity.translation_key == "voltage",
    )
    .replaces(AqaraW500ManufacturerCluster)
    .enum(
        AqaraW500ManufacturerCluster.AttributeDefs.sensor_type.name,
        AqaraW500SensorType,
        AqaraW500ManufacturerCluster.cluster_id,
        translation_key="sensor_type",
        fallback_name="Sensor type",
    )
    .enum(
        AqaraW500ManufacturerCluster.AttributeDefs.ntc_type.name,
        AqaraW500NtcType,
        AqaraW500ManufacturerCluster.cluster_id,
        translation_key="ntc_type",
        fallback_name="NTC type",
    )
    .enum(
        AqaraW500ManufacturerCluster.AttributeDefs.preset_mode.name,
        AqaraW500PresetMode,
        AqaraW500ManufacturerCluster.cluster_id,
        translation_key="preset_mode",
        fallback_name="Preset mode",
    )
    .switch(
        AqaraW500ManufacturerCluster.AttributeDefs.child_lock.name,
        AqaraW500ManufacturerCluster.cluster_id,
        off_value=0,
        on_value=1,
        translation_key="child_lock",
        fallback_name="Child lock",
    )
    .switch(
        AqaraW500ManufacturerCluster.AttributeDefs.window_detection.name,
        AqaraW500ManufacturerCluster.cluster_id,
        off_value=0,
        on_value=1,
        translation_key="window_detection",
        fallback_name="Window detection",
    )
    .number(
        AqaraW500ManufacturerCluster.AttributeDefs.hysteresis.name,
        AqaraW500ManufacturerCluster.cluster_id,
        min_value=0,
        max_value=3,
        step=0.5,
        unit=UnitOfTemperature.CELSIUS,
        multiplier=0.1,
        translation_key="hysteresis",
        fallback_name="Hysteresis",
    )
    .number(
        AqaraW500ManufacturerCluster.AttributeDefs.home_preset_temperature.name,
        AqaraW500ManufacturerCluster.cluster_id,
        min_value=5,
        max_value=35,
        step=0.5,
        unit=UnitOfTemperature.CELSIUS,
        multiplier=0.01,
        translation_key="home_preset_temperature",
        fallback_name="Home preset temperature",
    )
    .number(
        AqaraW500ManufacturerCluster.AttributeDefs.away_preset_temperature.name,
        AqaraW500ManufacturerCluster.cluster_id,
        min_value=5,
        max_value=35,
        step=0.5,
        unit=UnitOfTemperature.CELSIUS,
        multiplier=0.01,
        translation_key="away_preset_temperature",
        fallback_name="Away preset temperature",
    )
    .number(
        AqaraW500ManufacturerCluster.AttributeDefs.sleep_preset_temperature.name,
        AqaraW500ManufacturerCluster.cluster_id,
        min_value=5,
        max_value=35,
        step=0.5,
        unit=UnitOfTemperature.CELSIUS,
        multiplier=0.01,
        translation_key="sleep_preset_temperature",
        fallback_name="Sleep preset temperature",
    )
    .number(
        AqaraW500ManufacturerCluster.AttributeDefs.vacation_preset_temperature.name,
        AqaraW500ManufacturerCluster.cluster_id,
        min_value=5,
        max_value=35,
        step=0.5,
        unit=UnitOfTemperature.CELSIUS,
        multiplier=0.01,
        translation_key="vacation_preset_temperature",
        fallback_name="Vacation preset temperature",
    )
    .number(
        AqaraW500ManufacturerCluster.AttributeDefs.evening_preset_temperature.name,
        AqaraW500ManufacturerCluster.cluster_id,
        min_value=5,
        max_value=35,
        step=0.5,
        unit=UnitOfTemperature.CELSIUS,
        multiplier=0.01,
        translation_key="evening_preset_temperature",
        fallback_name="Wind down preset temperature",
    )
    .replaces(AqaraW500ThermostatCluster)
    .prevent_default_entity_creation(
        endpoint_id=1,
        cluster_id=AqaraW500ThermostatCluster.cluster_id,
        function=lambda entity: entity.translation_key == "min_heat_setpoint_limit",
    )
    .prevent_default_entity_creation(
        endpoint_id=1,
        cluster_id=AqaraW500ThermostatCluster.cluster_id,
        function=lambda entity: entity.translation_key == "max_heat_setpoint_limit",
    )
    .sensor(
        AqaraW500ThermostatCluster.AttributeDefs.local_temperature.name,
        AqaraW500ThermostatCluster.cluster_id,
        device_class=SensorDeviceClass.TEMPERATURE,
        state_class=SensorStateClass.MEASUREMENT,
        unit=UnitOfTemperature.CELSIUS,
        divisor=100,
        translation_key="temperature",
        fallback_name="Temperature",
    )
    .add_to_registry()
)