esphome-lsm6ds/esphome/components/zigbee/zigbee_zephyr.py

from datetime import datetime

from esphome import automation
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_NAME, __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_zephyr import (
    CONF_ON_JOIN,
    CONF_WIPE_ON_BOOT,
    CONF_ZIGBEE_BINARY_SENSOR,
    CONF_ZIGBEE_ID,
    KEY_EP_NUMBER,
    KEY_ZIGBEE,
    ZB_ZCL_BASIC_ATTRS_EXT_T,
    ZB_ZCL_CLUSTER_ID_BASIC,
    ZB_ZCL_CLUSTER_ID_BINARY_INPUT,
    ZB_ZCL_CLUSTER_ID_IDENTIFY,
    ZB_ZCL_IDENTIFY_ATTRS_T,
    BinaryAttrs,
    ZigbeeComponent,
    zigbee_ns,
)

ZigbeeBinarySensor = zigbee_ns.class_("ZigbeeBinarySensor", 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
        ),
    }
)


async def zephyr_to_code(config: ConfigType) -> None:
    zephyr_add_prj_conf("ZIGBEE", True)
    zephyr_add_prj_conf("ZIGBEE_APP_UTILS", True)
    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)

    if config[CONF_WIPE_ON_BOOT]:
        cg.add_define("USE_ZIGBEE_WIPE_ON_BOOT")
    var = cg.new_Pvariable(config[CONF_ID])

    if on_join_config := config.get(CONF_ON_JOIN):
        await automation.build_automation(var.get_join_trigger(), [], on_join_config)

    await cg.register_component(var, config)

    await _attr_to_code(config)
    CORE.add_job(_ctx_to_code, config)


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, CORE.name),
        zigbee_set_string(
            basic_attrs.date_code, datetime.now().strftime("%d/%m/%y %H:%M")
        ),
        zigbee_assign(
            basic_attrs.power_source,
            cg.RawExpression("ZB_ZCL_BASIC_POWER_SOURCE_DC_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)."""
    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,
) -> 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}, {', '.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 _add_binary_sensor(entity: cg.MockObj, config: ConfigType) -> None:
    # Find the next available endpoint slot
    slot_index = next(
        (i for i, v in enumerate(CORE.data[KEY_ZIGBEE][KEY_EP_NUMBER]) if v == ""), None
    )

    # Create unique names for this sensor's variables based on slot index
    prefix = f"zigbee_ep{slot_index + 1}"
    attrs_name = f"{prefix}_binary_attrs"
    attr_list_name = f"{prefix}_binary_input_attrib_list"
    cluster_list_name = f"{prefix}_cluster_list"
    ep_name = f"{prefix}_ep"

    # Create the binary attributes structure
    binary_attrs = zigbee_new_variable(attrs_name, BinaryAttrs)
    attr_list = zigbee_new_attr_list(
        attr_list_name,
        "ESPHOME_ZB_ZCL_DECLARE_BINARY_INPUT_ATTRIB_LIST",
        zigbee_assign(binary_attrs.out_of_service, 0),
        zigbee_assign(binary_attrs.present_value, 0),
        zigbee_assign(binary_attrs.status_flags, 0),
        zigbee_set_string(binary_attrs.description, config[CONF_NAME]),
    )

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

    # Create the ZigbeeBinarySensor component
    var = cg.new_Pvariable(config[CONF_ZIGBEE_BINARY_SENSOR], entity)
    await cg.register_component(var, config)

    cg.add(var.set_end_point(slot_index + 1))
    cg.add(var.set_cluster_attributes(binary_attrs))
    hub = await cg.get_variable(config[CONF_ZIGBEE_ID])
    cg.add(var.set_parent(hub))