Zemismart_ZTB25.py

import asyncio
import logging

from zigpy.profiles import zha
from zigpy.device import Device
import zigpy.types as t
from zigpy.quirks import CustomCluster

from zigpy.zcl.clusters.general import (
    Basic,
    GreenPowerProxy,
    Groups,
    Identify,
    OnOff,
    Ota,
    Scenes,
    Time,
)

from zhaquirks.const import (
    DEVICE_TYPE,
    ENDPOINTS,
    INPUT_CLUSTERS,
#    MODEL,
    OUTPUT_CLUSTERS,
    PROFILE_ID,
    MODELS_INFO,
)
from zhaquirks.tuya import (
    TuyaZBE000Cluster,
    TuyaZBExternalSwitchTypeCluster,
    TuyaZBOnOffAttributeCluster,
    TuyaOnOff
)
from zhaquirks.tuya.ts000x import Switch_3G_GPP

_LOGGER = logging.getLogger(__name__)

async def cast_tuya_magic_spell_task(
    dev: Device, tries: int = 100, rejoin: bool = False
) -> None:
    """Initialize device so that all endpoints become available."""
    import inspect

    basic_cluster = dev.endpoints[1].in_clusters[0]

    # The magic spell is needed only once.
    # TODO: Improve by doing this only once (successfully).

    # Magic spell - part 1
    # Note: attribute order is important
    attr_to_read = [
        "manufacturer",
        "zcl_version",
        "app_version",
        "model",
        "power_source",
        0xFFFE,
    ]
    if "tries" in inspect.getfullargspec(basic_cluster.read_attributes)[0]:
        _LOGGER.debug(f"Cast Tuya Magic Spell on {dev.ieee!r} with {tries} tries")
        res = await basic_cluster.read_attributes(attr_to_read, tries=tries)
    else:
        _LOGGER.debug(f"Cast Tuya Magic Spell on {dev.ieee!r}")
        res = await basic_cluster.read_attributes(attr_to_read)

    _LOGGER.debug(f"Tuya Magic Spell result {res!r} for {dev.ieee!r}")

    # Magic spell - part 2 (skipped - does not seem to be needed)
    # attr_to_write={0xffde:13}
    # basic_cluster.write_attributes(attr_to_write)

    if rejoin:
        # Leave with rejoin - may need to be adjuste to work everywhere
        # or require a minimum zigpy version
        # This should have the device leave and rejoin immediately triggering
        # the discovery of the endpoints that appear after the magic trick

        # Note: this is not validated yet and disabled by default
        _LOGGER.debug(f"Send leave with rejoin request to {dev.ieee!r}")
        res = await dev.zdo.request(0x0034, dev.ieee, 0x01, tries)
        _LOGGER.debug(f"Leave with rejoin result {res!r} for {dev.ieee!r}")

        app = dev.application
        # Delete the device from the database
        app.listener_event("device_removed", dev)

        # Delete the device from zigpy
        app.devices.pop(dev.ieee, None)


def cast_tuya_magic_spell(dev: Device, tries: int = 3) -> None:
    """Set up the magic spell asynchronously."""

    # Note for sleepy devices the number of tries may need to be increased to 100.

    dev._magic_spell_task = asyncio.create_task(
        cast_tuya_magic_spell_task(dev, tries=tries)
    )


class TuyaBasicCluster(CustomCluster, Basic):
    """Provide Tuya Basic Cluster with magic spell."""

    attributes = Basic.attributes.copy()
    attributes.update(
        {
            0xFFDE: ("tuya_FFDE", t.uint8_t, True),
            # 0xffe0: ("tuya_FFE0", TODO.Array, True),
            # 0xffe1: ("tuya_FFE1", TODO.Array, True),
            0xFFE2: ("tuya_FFE2", t.uint8_t, True),
            # 0xffe3: ("tuya_FFE3", TODO.Array, True),
        }
    )

    async def bind(self):
        """Bind cluster."""

        _LOGGER.debug(
            f"Requesting Tuya Magic Spell for {self.ieee!r} in basic bind method"
        )
        tries = 3
        await asyncio.create_task(cast_tuya_magic_spell_task(self, tries=tries))

        return await super().bind()

class Switch_3G_Zemismart_ZTB25(Switch_3G_GPP):

    def __init__(self, *args, **kwargs):
        self.signature[MODELS_INFO] = [("_TZ3000_vjhcenzo", "TS0003")]
        self.replacement[ENDPOINTS][3][DEVICE_TYPE] = zha.DeviceType.ON_OFF_PLUG_IN_UNIT
        super().__init__(*args, **kwargs)
        cast_tuya_magic_spell(self, tries=3)