esphome/esphome/components/zigbee/zigbee_zephyr.py

import datetime
import random

import esphome.codegen as cg
from esphome.components.zephyr import zephyr_add_prj_conf
import esphome.config_validation as cv
from esphome.const import (
    CONF_ID,
    CONF_MODEL,
    CONF_NAME,
    CONF_UNIT_OF_MEASUREMENT,
    __version__,
)
from esphome.core import CORE, CoroPriority, coroutine_with_priority
from esphome.cpp_generator import (
    AssignmentExpression,
    MockObj,
    VariableDeclarationExpression,
)
from esphome.types import ConfigType

from .const import (
    BACNET_UNIT_NO_UNITS,
    BACNET_UNITS,
    CONF_POWER_SOURCE,
    CONF_ROUTER,
    CONF_WIPE_ON_BOOT,
    KEY_ZIGBEE,
    POWER_SOURCE,
    AnalogAttrs,
    AnalogAttrsOutput,
    BinaryAttrs,
    ZigbeeComponent,
    zigbee_ns,
)
from .const_zephyr import (
    CONF_IEEE802154_VENDOR_OUI,
    CONF_SLEEPY,
    CONF_ZIGBEE_BINARY_SENSOR,
    CONF_ZIGBEE_ID,
    CONF_ZIGBEE_NUMBER,
    CONF_ZIGBEE_SENSOR,
    CONF_ZIGBEE_SWITCH,
    KEY_EP_NUMBER,
    ZB_ZCL_BASIC_ATTRS_EXT_T,
    ZB_ZCL_CLUSTER_ID_ANALOG_INPUT,
    ZB_ZCL_CLUSTER_ID_ANALOG_OUTPUT,
    ZB_ZCL_CLUSTER_ID_BASIC,
    ZB_ZCL_CLUSTER_ID_BINARY_INPUT,
    ZB_ZCL_CLUSTER_ID_BINARY_OUTPUT,
    ZB_ZCL_CLUSTER_ID_IDENTIFY,
    ZB_ZCL_IDENTIFY_ATTRS_T,
)

ZigbeeBinarySensor = zigbee_ns.class_("ZigbeeBinarySensor", cg.Component)
ZigbeeSensor = zigbee_ns.class_("ZigbeeSensor", cg.Component)
ZigbeeSwitch = zigbee_ns.class_("ZigbeeSwitch", cg.Component)
ZigbeeNumber = zigbee_ns.class_("ZigbeeNumber", cg.Component)

zephyr_binary_sensor = cv.Schema(
    {
        cv.OnlyWith(CONF_ZIGBEE_ID, ["nrf52", "zigbee"]): cv.use_id(ZigbeeComponent),
        cv.OnlyWith(CONF_ZIGBEE_BINARY_SENSOR, ["nrf52", "zigbee"]): cv.declare_id(
            ZigbeeBinarySensor
        ),
    }
)

zephyr_sensor = cv.Schema(
    {
        cv.OnlyWith(CONF_ZIGBEE_ID, ["nrf52", "zigbee"]): cv.use_id(ZigbeeComponent),
        cv.OnlyWith(CONF_ZIGBEE_SENSOR, ["nrf52", "zigbee"]): cv.declare_id(
            ZigbeeSensor
        ),
    }
)

zephyr_switch = cv.Schema(
    {
        cv.OnlyWith(CONF_ZIGBEE_ID, ["nrf52", "zigbee"]): cv.use_id(ZigbeeComponent),
        cv.OnlyWith(CONF_ZIGBEE_SWITCH, ["nrf52", "zigbee"]): cv.declare_id(
            ZigbeeSwitch
        ),
    }
)

zephyr_number = cv.Schema(
    {
        cv.OnlyWith(CONF_ZIGBEE_ID, ["nrf52", "zigbee"]): cv.use_id(ZigbeeComponent),
        cv.OnlyWith(CONF_ZIGBEE_NUMBER, ["nrf52", "zigbee"]): cv.declare_id(
            ZigbeeNumber
        ),
    }
)


async def zephyr_to_code(config: ConfigType) -> "MockObj":
    zephyr_add_prj_conf("ZIGBEE", True)
    zephyr_add_prj_conf("ZIGBEE_APP_UTILS", True)
    if config[CONF_ROUTER]:
        zephyr_add_prj_conf("ZIGBEE_ROLE_ROUTER", True)
    else:
        zephyr_add_prj_conf("ZIGBEE_ROLE_END_DEVICE", True)

    zephyr_add_prj_conf("ZIGBEE_CHANNEL_SELECTION_MODE_MULTI", True)

    zephyr_add_prj_conf("CRYPTO", True)

    zephyr_add_prj_conf("NET_IPV6", False)
    zephyr_add_prj_conf("NET_IP_ADDR_CHECK", False)
    zephyr_add_prj_conf("NET_UDP", False)

    # disable all extra to reduce power and save flash
    zephyr_add_prj_conf("ZIGBEE_HAVE_SERIAL", False)
    zephyr_add_prj_conf("ZBOSS_ERROR_PRINT_TO_LOG", False)
    zephyr_add_prj_conf("DK_LIBRARY", False)

    cg.add_build_flag("-Wl,--wrap=zb_zcl_put_reporting_info_from_req")

    # Wrap the transceiver sleep/receive/transmit calls to measure how long the
    # radio is powered down. The span between a zb_trans_enter_sleep() and the
    # following zb_trans_enter_receive() or zb_trans_transmit() is time the
    # radio spent asleep.
    cg.add_build_flag("-Wl,--wrap=zb_trans_enter_sleep")
    cg.add_build_flag("-Wl,--wrap=zb_trans_enter_receive")
    cg.add_build_flag("-Wl,--wrap=zb_trans_transmit")

    if CONF_IEEE802154_VENDOR_OUI in config:
        zephyr_add_prj_conf("IEEE802154_VENDOR_OUI_ENABLE", True)
        random_number = config[CONF_IEEE802154_VENDOR_OUI]
        if random_number == "random":
            random_number = random.randint(0x000000, 0xFFFFFF)
        zephyr_add_prj_conf("IEEE802154_VENDOR_OUI", random_number)

    if config[CONF_WIPE_ON_BOOT]:
        if config[CONF_WIPE_ON_BOOT] == "once":
            cg.add_define(
                "USE_ZIGBEE_WIPE_ON_BOOT_MAGIC", random.randint(0x000001, 0xFFFFFF)
            )
        cg.add_define("USE_ZIGBEE_WIPE_ON_BOOT")

    # Generate attribute lists before any await that could yield (e.g., build_automation
    # waiting for variables from other components). If the hub's priority decays while
    # yielding, deferred entity jobs may add cluster list globals that reference these
    # attribute lists before they're declared.
    await _attr_to_code(config)

    var = cg.new_Pvariable(config[CONF_ID])

    await cg.register_component(var, config)

    CORE.add_job(_ctx_to_code, config)

    cg.add(var.set_sleepy(config[CONF_SLEEPY]))

    return var


async def _attr_to_code(config: ConfigType) -> None:
    # Create the basic attributes structure and attribute list
    basic_attrs = zigbee_new_variable("zigbee_basic_attrs", ZB_ZCL_BASIC_ATTRS_EXT_T)
    zigbee_new_attr_list(
        "zigbee_basic_attrib_list",
        "ZB_ZCL_DECLARE_BASIC_ATTRIB_LIST_EXT",
        zigbee_assign(basic_attrs.zcl_version, cg.RawExpression("ZB_ZCL_VERSION")),
        zigbee_assign(basic_attrs.app_version, 0),
        zigbee_assign(basic_attrs.stack_version, 0),
        zigbee_assign(basic_attrs.hw_version, 0),
        zigbee_set_string(basic_attrs.mf_name, "esphome"),
        zigbee_set_string(basic_attrs.model_id, config[CONF_MODEL]),
        zigbee_set_string(
            basic_attrs.date_code,
            # Local build time, matching the esp32 implementation
            # (App.get_build_time() in C++).
            datetime.datetime.now().astimezone().strftime("%Y%m%d %H%M%S"),
        ),
        zigbee_assign(
            basic_attrs.power_source,
            POWER_SOURCE[config[CONF_POWER_SOURCE]],
        ),
        zigbee_set_string(basic_attrs.location_id, ""),
        zigbee_assign(
            basic_attrs.ph_env, cg.RawExpression("ZB_ZCL_BASIC_ENV_UNSPECIFIED")
        ),
        zigbee_set_string(basic_attrs.sw_ver, __version__),
    )

    # Create the identify attributes structure and attribute list
    identify_attrs = zigbee_new_variable(
        "zigbee_identify_attrs", ZB_ZCL_IDENTIFY_ATTRS_T
    )
    zigbee_new_attr_list(
        "zigbee_identify_attrib_list",
        "ZB_ZCL_DECLARE_IDENTIFY_ATTRIB_LIST",
        zigbee_assign(
            identify_attrs.identify_time,
            cg.RawExpression("ZB_ZCL_IDENTIFY_IDENTIFY_TIME_DEFAULT_VALUE"),
        ),
    )


def zigbee_new_variable(name: str, type_: str) -> cg.MockObj:
    """Create a global variable with the given name and type."""
    decl = VariableDeclarationExpression(type_, "", name)
    CORE.add_global(decl)
    return MockObj(name, ".")


def zigbee_assign(target: cg.MockObj, expression: cg.RawExpression | int) -> str:
    """Assign an expression to a target and return a reference to it."""
    cg.add(AssignmentExpression("", "", target, expression))
    return f"&{target}"


def zigbee_set_string(target: cg.MockObj, value: str) -> str:
    """Set a ZCL string value and return the target name (arrays decay to pointers)."""
    # Zigbee supports only ASCII
    value = value.encode("ascii", "ignore").decode()
    cg.add(
        cg.RawExpression(
            f"ZB_ZCL_SET_STRING_VAL({target}, {cg.safe_exp(value)}, ZB_ZCL_STRING_CONST_SIZE({cg.safe_exp(value)}))"
        )
    )
    return str(target)


def zigbee_new_attr_list(name: str, macro: str, *args: str) -> str:
    """Create an attribute list using a ZBOSS macro and return the name."""
    obj = cg.RawExpression(f"{macro}({name}, {', '.join(args)})")
    CORE.add_global(obj)
    return name


class ZigbeeClusterDesc:
    """Represents a Zigbee cluster descriptor for code generation."""

    def __init__(self, cluster_id: str, attr_list_name: str | None = None) -> None:
        self._cluster_id = cluster_id
        self._attr_list_name = attr_list_name

    @property
    def cluster_id(self) -> str:
        return self._cluster_id

    @property
    def has_attrs(self) -> bool:
        return self._attr_list_name is not None

    def __str__(self) -> str:
        role = (
            "ZB_ZCL_CLUSTER_SERVER_ROLE"
            if self._attr_list_name
            else "ZB_ZCL_CLUSTER_CLIENT_ROLE"
        )
        if self._attr_list_name:
            attr_count = f"ZB_ZCL_ARRAY_SIZE({self._attr_list_name}, zb_zcl_attr_t)"
            return f"ZB_ZCL_CLUSTER_DESC({self._cluster_id}, {attr_count}, {self._attr_list_name}, {role}, ZB_ZCL_MANUF_CODE_INVALID)"
        return f"ZB_ZCL_CLUSTER_DESC({self._cluster_id}, 0, NULL, {role}, ZB_ZCL_MANUF_CODE_INVALID)"


def zigbee_new_cluster_list(
    name: str, clusters: list[ZigbeeClusterDesc]
) -> tuple[str, list[ZigbeeClusterDesc]]:
    """Create a cluster list array and return its name and the clusters."""
    # Always include basic and identify clusters first
    all_clusters = [
        ZigbeeClusterDesc(ZB_ZCL_CLUSTER_ID_BASIC, "zigbee_basic_attrib_list"),
        ZigbeeClusterDesc(ZB_ZCL_CLUSTER_ID_IDENTIFY, "zigbee_identify_attrib_list"),
    ]
    all_clusters.extend(clusters)

    cluster_strs = [str(c) for c in all_clusters]
    CORE.add_global(
        cg.RawExpression(
            f"zb_zcl_cluster_desc_t {name}[] = {{{', '.join(cluster_strs)}}}"
        )
    )
    return (name, all_clusters)


def zigbee_register_ep(
    ep_name: str,
    cluster_list_name: str,
    report_attr_count: int,
    clusters: list[ZigbeeClusterDesc],
    slot_index: int,
    app_device_id: str,
) -> None:
    """Register a Zigbee endpoint."""
    in_cluster_num = sum(1 for c in clusters if c.has_attrs)
    out_cluster_num = len(clusters) - in_cluster_num
    cluster_ids = [c.cluster_id for c in clusters]

    # Store endpoint name for device context generation
    CORE.data[KEY_ZIGBEE][KEY_EP_NUMBER][slot_index] = ep_name

    # Generate the endpoint declaration
    ep_id = slot_index + 1  # Endpoints are 1-indexed
    obj = cg.RawExpression(
        f"ESPHOME_ZB_HA_DECLARE_EP({ep_name}, {ep_id}, {cluster_list_name}, "
        f"{in_cluster_num}, {out_cluster_num}, {report_attr_count}, {app_device_id}, {', '.join(cluster_ids)})"
    )
    CORE.add_global(obj)


@coroutine_with_priority(CoroPriority.LATE)
async def _ctx_to_code(config: ConfigType) -> None:
    cg.add_define("ZIGBEE_ENDPOINTS_COUNT", len(CORE.data[KEY_ZIGBEE][KEY_EP_NUMBER]))
    cg.add_global(
        cg.RawExpression(
            f"ZBOSS_DECLARE_DEVICE_CTX_EP_VA(zb_device_ctx, &{', &'.join(CORE.data[KEY_ZIGBEE][KEY_EP_NUMBER])})"
        )
    )
    cg.add(cg.RawExpression("ZB_AF_REGISTER_DEVICE_CTX(&zb_device_ctx)"))


async def zephyr_setup_binary_sensor(entity: cg.MockObj, config: ConfigType) -> None:
    CORE.add_job(_add_binary_sensor, entity, config)


async def zephyr_setup_sensor(entity: cg.MockObj, config: ConfigType) -> None:
    CORE.add_job(_add_sensor, entity, config)


async def zephyr_setup_switch(entity: cg.MockObj, config: ConfigType) -> None:
    CORE.add_job(_add_switch, entity, config)


async def zephyr_setup_number(
    entity: cg.MockObj,
    config: ConfigType,
    min_value: float,
    max_value: float,
    step: float,
) -> None:
    CORE.add_job(_add_number, entity, config, min_value, max_value, step)


def get_slot_index() -> int:
    """Find the next available endpoint slot."""
    slot = next(
        (i for i, v in enumerate(CORE.data[KEY_ZIGBEE][KEY_EP_NUMBER]) if v == ""), None
    )
    if slot is None:
        raise cv.Invalid(
            f"No available Zigbee endpoint slots ({len(CORE.data[KEY_ZIGBEE][KEY_EP_NUMBER])} in use)"
        )
    return slot


async def _add_zigbee_ep(
    entity: cg.MockObj,
    config: ConfigType,
    component_key,
    attrs_type,
    zcl_macro: str,
    cluster_id: str,
    app_device_id: str,
    extra_field_values: dict[str, int] | None = None,
) -> None:
    slot_index = get_slot_index()

    prefix = f"zigbee_ep{slot_index + 1}"
    attrs_name = f"{prefix}_attrs"
    attr_list_name = f"{prefix}_attrib_list"
    cluster_list_name = f"{prefix}_cluster_list"
    ep_name = f"{prefix}_ep"

    # Create attribute struct
    attrs = zigbee_new_variable(attrs_name, attrs_type)

    # Build attribute list args
    attr_args = [
        zigbee_assign(attrs.out_of_service, 0),
        zigbee_assign(attrs.present_value, 0),
        zigbee_assign(attrs.status_flags, 0),
    ]
    # Add extra field assignments (e.g., engineering_units for sensors)
    if extra_field_values:
        for field_name, value in extra_field_values.items():
            attr_args.append(zigbee_assign(getattr(attrs, field_name), value))
    attr_args.append(zigbee_set_string(attrs.description, config[CONF_NAME]))

    # Create attribute list
    attr_list = zigbee_new_attr_list(attr_list_name, zcl_macro, *attr_args)

    # Create cluster list and register endpoint
    cluster_list_name, clusters = zigbee_new_cluster_list(
        cluster_list_name,
        [ZigbeeClusterDesc(cluster_id, attr_list)],
    )
    zigbee_register_ep(
        ep_name, cluster_list_name, 2, clusters, slot_index, app_device_id
    )

    # Create ESPHome component
    var = cg.new_Pvariable(config[component_key], entity)
    await cg.register_component(var, {})

    cg.add(var.set_endpoint(slot_index + 1))
    cg.add(var.set_cluster_attributes(attrs))

    hub = await cg.get_variable(config[CONF_ZIGBEE_ID])
    cg.add(var.set_parent(hub))


async def _add_binary_sensor(entity: cg.MockObj, config: ConfigType) -> None:
    await _add_zigbee_ep(
        entity,
        config,
        CONF_ZIGBEE_BINARY_SENSOR,
        BinaryAttrs,
        "ESPHOME_ZB_ZCL_DECLARE_BINARY_INPUT_ATTRIB_LIST",
        ZB_ZCL_CLUSTER_ID_BINARY_INPUT,
        "ZB_HA_SIMPLE_SENSOR_DEVICE_ID",
    )


async def _add_sensor(entity: cg.MockObj, config: ConfigType) -> None:
    # Get BACnet engineering unit from unit_of_measurement
    unit = config.get(CONF_UNIT_OF_MEASUREMENT, "")
    bacnet_unit = BACNET_UNITS.get(unit, BACNET_UNIT_NO_UNITS)

    await _add_zigbee_ep(
        entity,
        config,
        CONF_ZIGBEE_SENSOR,
        AnalogAttrs,
        "ESPHOME_ZB_ZCL_DECLARE_ANALOG_INPUT_ATTRIB_LIST",
        ZB_ZCL_CLUSTER_ID_ANALOG_INPUT,
        "ZB_HA_CUSTOM_ATTR_DEVICE_ID",
        extra_field_values={"engineering_units": bacnet_unit},
    )


async def _add_switch(entity: cg.MockObj, config: ConfigType) -> None:
    await _add_zigbee_ep(
        entity,
        config,
        CONF_ZIGBEE_SWITCH,
        BinaryAttrs,
        "ESPHOME_ZB_ZCL_DECLARE_BINARY_OUTPUT_ATTRIB_LIST",
        ZB_ZCL_CLUSTER_ID_BINARY_OUTPUT,
        "ZB_HA_CUSTOM_ATTR_DEVICE_ID",
    )


async def _add_number(
    entity: cg.MockObj,
    config: ConfigType,
    min_value: float,
    max_value: float,
    step: float,
) -> None:
    # Get BACnet engineering unit from unit_of_measurement
    unit = config.get(CONF_UNIT_OF_MEASUREMENT, "")
    bacnet_unit = BACNET_UNITS.get(unit, BACNET_UNIT_NO_UNITS)

    await _add_zigbee_ep(
        entity,
        config,
        CONF_ZIGBEE_NUMBER,
        AnalogAttrsOutput,
        "ESPHOME_ZB_ZCL_DECLARE_ANALOG_OUTPUT_ATTRIB_LIST",
        ZB_ZCL_CLUSTER_ID_ANALOG_OUTPUT,
        "ZB_HA_CUSTOM_ATTR_DEVICE_ID",
        extra_field_values={
            "max_present_value": max_value,
            "min_present_value": min_value,
            "resolution": step,
            "engineering_units": bacnet_unit,
        },
    )