shelly_gen4_1pm.py

"""Shelly 1PM Gen4 Pure V2 Quirk using the QuirkBuilder pattern.

This version uses the correct AttributeDefs pattern within the custom clusters
to ensure the .replaces() method works correctly during initialization.
"""

from __future__ import annotations

import enum
import logging

# Only QuirkBuilder and EntityType are needed from the v2 module
from zigpy.quirks.v2 import QuirkBuilder, EntityType
from zigpy.quirks import CustomCluster
import zigpy.types as t
from zigpy.zcl.foundation import BaseAttributeDefs, ZCLAttributeDef


_LOGGER = logging.getLogger(__name__)

# #############################################################################
#
#   Part 1: Shared Manufacturer Logic
#
# #############################################################################

SHELLY_MANUFACTURER = "Shelly"
SHELLY_MANUFACTURER_ID = 0x1490
RPC_CLUSTER_ID = 0xFC01
WIFI_SETUP_CLUSTER_ID = 0xFC02
SHELLY_EP = 239


class WifiAction(t.enum8):
    """Actions for the WiFi Setup cluster."""

    NOP = 0
    APPLY = 1
    RESET = 2


class ShellyRPCCluster(CustomCluster):
    """Shelly RPC custom cluster (0xFC01)."""

    cluster_id = RPC_CLUSTER_ID
    name = "Shelly RPC"
    ep_attribute = "shelly_rpc"
    manufacturer_id_override = SHELLY_MANUFACTURER_ID

    class AttributeDefs(BaseAttributeDefs):
        """Attribute definitions."""

        data = ZCLAttributeDef(id=0x0000, type=bytes, is_manufacturer_specific=True)
        tx_ctl = ZCLAttributeDef(id=0x0001, type=t.uint16_t, is_manufacturer_specific=True)
        rx_ctl = ZCLAttributeDef(id=0x0002, type=t.uint16_t, is_manufacturer_specific=True)


class ShellyWiFiSetupCluster(CustomCluster):
    """Shelly WiFi Setup custom cluster (0xFC02)."""

    cluster_id = WIFI_SETUP_CLUSTER_ID
    name = "Shelly WiFi Setup"
    ep_attribute = "shelly_wifi_setup"
    manufacturer_id_override = SHELLY_MANUFACTURER_ID

    class AttributeDefs(BaseAttributeDefs):
        """Attribute definitions."""

        status = ZCLAttributeDef(id=0x0000, type=str, is_manufacturer_specific=True)
        ip = ZCLAttributeDef(id=0x0001, type=str, is_manufacturer_specific=True)
        action = ZCLAttributeDef(id=0x0002, type=WifiAction, is_manufacturer_specific=True)
        dhcp = ZCLAttributeDef(id=0x0003, type=bool, is_manufacturer_specific=True)
        enable = ZCLAttributeDef(id=0x0004, type=bool, is_manufacturer_specific=True)
        ssid = ZCLAttributeDef(id=0x0005, type=str, is_manufacturer_specific=True)
        password = ZCLAttributeDef(id=0x0006, type=str, is_manufacturer_specific=True)


# #############################################################################
#
#   Part 2: Device-Specific Quirk Definition
#
# #############################################################################

(
    QuirkBuilder(SHELLY_MANUFACTURER, "1PM")
    .replaces(ShellyRPCCluster, endpoint_id=SHELLY_EP)
    .replaces(ShellyWiFiSetupCluster, endpoint_id=SHELLY_EP)
    # WiFi Setup Entities
    .sensor(
        attribute_name=ShellyWiFiSetupCluster.AttributeDefs.status.name,
        cluster_id=WIFI_SETUP_CLUSTER_ID,
        endpoint_id=SHELLY_EP,
        translation_key="wifi_status",
        fallback_name="WiFi Status",
        entity_type=EntityType.DIAGNOSTIC,
    )
    .sensor(
        attribute_name=ShellyWiFiSetupCluster.AttributeDefs.ip.name,
        cluster_id=WIFI_SETUP_CLUSTER_ID,
        endpoint_id=SHELLY_EP,
        translation_key="wifi_ip",
        fallback_name="WiFi IP",
        entity_type=EntityType.DIAGNOSTIC,
    )
    .switch(
        attribute_name=ShellyWiFiSetupCluster.AttributeDefs.enable.name,
        cluster_id=WIFI_SETUP_CLUSTER_ID,
        endpoint_id=SHELLY_EP,
        translation_key="wifi_enabled",
        fallback_name="WiFi Enabled",
        entity_type=EntityType.CONFIG,
    )
    .switch(
        attribute_name=ShellyWiFiSetupCluster.AttributeDefs.dhcp.name,
        cluster_id=WIFI_SETUP_CLUSTER_ID,
        endpoint_id=SHELLY_EP,
        translation_key="wifi_dhcp",
        fallback_name="WiFi DHCP",
        entity_type=EntityType.CONFIG,
    )
    .enum(
        attribute_name=ShellyWiFiSetupCluster.AttributeDefs.action.name,
        enum_class=WifiAction,
        cluster_id=WIFI_SETUP_CLUSTER_ID,
        endpoint_id=SHELLY_EP,
        translation_key="wifi_action",
        fallback_name="WiFi Action",
        entity_type=EntityType.CONFIG,
    )
    # RPC Entities
    .number(
        attribute_name=ShellyRPCCluster.AttributeDefs.tx_ctl.name,
        cluster_id=RPC_CLUSTER_ID,
        endpoint_id=SHELLY_EP,
        translation_key="rpc_tx_control",
        fallback_name="RPC Tx Control",
        entity_type=EntityType.CONFIG,
    )
    .sensor(
        attribute_name=ShellyRPCCluster.AttributeDefs.rx_ctl.name,
        cluster_id=RPC_CLUSTER_ID,
        endpoint_id=SHELLY_EP,
        translation_key="rpc_rx_control",
        fallback_name="RPC Rx Control",
        entity_type=EntityType.DIAGNOSTIC,
    )
    .add_to_registry()
)