location_managemet/app/drivers/tuya_cloud/driver.py

"""Tuya Cloud Gateway Driver.

Implements the Tuya Device Sharing SDK integration.  A single
"board" record represents an entire Tuya / Smart Life account.
Each discovered Tuya device is stored as a TuyaDevice record.

Auth flow (QR-code based):
  1.  User provides any short string as their "user_code".
  2.  Driver calls LoginControl().qr_code() → receives a QR token.
  3.  User scans the QR with the Tuya/Smart Life app.
  4.  Driver polls LoginControl().login_result() until success.
  5.  Tokens (access_token, refresh_token, etc.) are saved in board.config.

Device categories → kinds mapping follows the Tuya OpenAPI category definitions.
"""
from __future__ import annotations

import json
import logging
from typing import TYPE_CHECKING

from app.drivers.base import BoardDriver

if TYPE_CHECKING:
    from app.models.board import Board

logger = logging.getLogger(__name__)

# ── Tuya sharing app credentials (same as Home Assistant) ─────────────────────
TUYA_CLIENT_ID = "HA_3y9q4ak7g4ephrvke"
TUYA_SCHEMA    = "haauthorize"

# ── Category code → kind ──────────────────────────────────────────────────────
CATEGORY_KIND: dict[str, str] = {
    # Switches / sockets / power strips
    "kg": "switch", "cz": "switch", "pc": "switch", "tdq": "switch",
    "tgkg": "switch", "tgq": "switch", "dlq": "switch", "zndb": "switch",
    # Lights
    "dj": "light",  "dd": "light",  "xdd": "light", "fwd": "light",
    "dc": "light",  "tydj": "light", "gyd": "light",
    # Fans
    "fs": "fan",    "fsd": "fan",
    # Sensors
    "wsdcg": "sensor", "co2bj": "sensor", "mcs": "sensor",
    "pir": "sensor",   "ldcg": "sensor",  "pm2.5": "sensor",
    # Covers / blinds
    "cl": "cover",  "clkg": "cover",
}

KIND_ICON: dict[str, str] = {
    "switch": "bi-toggles",
    "light":  "bi-lightbulb-fill",
    "fan":    "bi-fan",
    "sensor": "bi-thermometer-half",
    "cover":  "bi-door-open",
}


def category_kind(cat: str) -> str:
    return CATEGORY_KIND.get(cat, "switch")


# ── DP value formatting ────────────────────────────────────────────────────────

# DP code substrings that indicate a temperature value stored as integer × 10
_TEMP_DP_SUBSTRINGS  = ("temp",)
# Substrings that *contain* "temp" but are NOT scaled temperature readings
_TEMP_DP_EXCLUDE     = ("unit_convert", "alarm")
# DP code substrings that carry a raw humidity percentage (0-100)
_HUMID_DP_SUBSTRINGS = ("humidity", "va_humidity", "hum_set", "hum_alarm")


def format_dp_value(dp_code: str, value, status: dict | None = None):
    """Return a human-readable representation of a Tuya DP value.

    Rules
    -----
    * **Temperature DPs** (any code containing ``"temp"`` but not
      ``"unit_convert"`` or ``"alarm"``): raw integer is divided by 10
      and the unit symbol (°C / °F) is appended based on the device's
      ``temp_unit_convert`` status key.
    * **Humidity DPs**: raw integer is shown with a ``%`` suffix.
    * **Booleans**: returned unchanged so the template can render badges.
    * Everything else: returned unchanged.
    """
    if isinstance(value, bool):
        return value  # template renders True/False as coloured badges

    code = dp_code.lower()

    # Temperature (Tuya encodes as integer × 10)
    if (
        any(sub in code for sub in _TEMP_DP_SUBSTRINGS)
        and not any(exc in code for exc in _TEMP_DP_EXCLUDE)
        and isinstance(value, (int, float))
    ):
        unit_raw = (status or {}).get("temp_unit_convert", "c")
        unit_sym = "°F" if str(unit_raw).lower() == "f" else "°C"
        return f"{value / 10:.1f} {unit_sym}"

    # Humidity (raw 0-100 %)
    if any(sub in code for sub in _HUMID_DP_SUBSTRINGS) and isinstance(value, (int, float)):
        return f"{value} %"

    return value  # pass through; template handles display


def detect_switch_dps(device) -> list[str]:
    """Return ordered list of writable boolean switch DP codes for a device.

    Looks at device.function (writable DPs) and device.status keys.
    Returns e.g. ['switch'], ['switch_1','switch_2'], etc.
    """
    funcs  = device.function if hasattr(device, "function") and device.function else {}
    status = device.status   if hasattr(device, "status")   and device.status   else {}
    candidates = set(funcs.keys()) | set(status.keys())

    dps: list[str] = []
    if "switch" in candidates:
        dps.append("switch")
    for i in range(1, 10):
        key = f"switch_{i}"
        if key in candidates:
            dps.append(key)

    # Fallback: treat "led_switch" / "fan_switch" as single channel
    if not dps:
        for fallback in ("led_switch", "fan_switch", "switch_led"):
            if fallback in candidates:
                dps.append(fallback)
                break

    return dps or ["switch"]


def build_manager(board: "Board"):
    """Construct a tuya_sharing.Manager from this board's stored config.

    The Manager handles token-refresh automatically.  Any refreshed token
    is written back to board.config so the next save() persists it.
    """
    try:
        from tuya_sharing import Manager, SharingTokenListener
    except ImportError:
        logger.error("tuya-device-sharing-sdk not installed; run: pip install tuya-device-sharing-sdk")
        return None

    cfg          = board.config
    token_info   = cfg.get("tuya_token_info")
    user_code    = cfg.get("tuya_user_code", "")
    terminal_id  = cfg.get("tuya_terminal_id", "")
    endpoint     = cfg.get("tuya_endpoint", "https://apigw.iotbing.com")

    if not token_info or not user_code:
        logger.debug("Board %d has no Tuya token — skipping manager creation", board.id)
        return None

    class _Listener(SharingTokenListener):
        def update_token(self, new_token_info: dict):
            # Merge refreshed token back; caller must db.session.commit() later
            cfg["tuya_token_info"] = new_token_info
            board.config = cfg

    try:
        return Manager(
            client_id=TUYA_CLIENT_ID,
            user_code=user_code,
            terminal_id=terminal_id,
            end_point=endpoint,
            token_response=token_info,
            listener=_Listener(),
        )
    except Exception as exc:
        logger.error("Cannot create Tuya Manager for board %d: %s", board.id, exc)
        return None


# ── Driver class ───────────────────────────────────────────────────────────────

class TuyaCloudDriver(BoardDriver):
    DRIVER_ID      = "tuya_cloud"
    DISPLAY_NAME   = "Tuya Cloud Gateway"
    DESCRIPTION    = "Control Tuya / Smart Life devices via the Tuya Device Sharing cloud API"
    DEFAULT_NUM_RELAYS = 0
    DEFAULT_NUM_INPUTS = 0

    # ── Abstract stubs (not used for gateway boards) ──────────────────────────
    def get_relay_status(self, board, relay_num):  return None
    def set_relay(self, board, relay_num, state):  return False
    def toggle_relay(self, board, relay_num):      return None
    def register_webhook(self, board, callback_url): return True

    # ── Core poll ─────────────────────────────────────────────────────────────
    def poll(self, board: "Board") -> dict:
        """Sync device list + states from Tuya cloud."""
        try:
            self.sync_devices(board)
            board.is_online = True
        except Exception as exc:
            logger.warning("Tuya poll failed for board %d: %s", board.id, exc)
            board.is_online = False
        return {"relay_states": {}, "input_states": {}, "is_online": board.is_online}

    # ── Sync ──────────────────────────────────────────────────────────────────
    def sync_devices(self, board: "Board") -> int:
        """Fetch the device list from Tuya cloud and upsert TuyaDevice rows."""
        from app import db
        from app.models.tuya_device import TuyaDevice
        from datetime import datetime

        mgr = build_manager(board)
        if mgr is None:
            return 0

        mgr.update_device_cache()

        # Persist any token that was refreshed during the update
        try:
            db.session.commit()
        except Exception:
            db.session.rollback()

        existing: dict[str, TuyaDevice] = {
            d.device_id: d
            for d in TuyaDevice.query.filter_by(board_id=board.id).all()
        }

        synced = 0
        for device_id, device in mgr.device_map.items():
            dev = existing.get(device_id)
            if dev is None:
                dev = TuyaDevice(board_id=board.id, device_id=device_id)
                db.session.add(dev)

            if not dev.name:
                dev.name = device.name or device_id
            dev.category     = device.category or ""
            dev.product_name = getattr(device, "product_name", "") or ""
            dev.is_online    = device.online
            dev.kind         = category_kind(dev.category)
            dev.kind_icon_cls = KIND_ICON.get(dev.kind, "bi-toggles")

            switch_dps = detect_switch_dps(device)
            dev.num_channels     = len(switch_dps)
            dev.switch_dps_json  = json.dumps(switch_dps)
            dev.status_json      = json.dumps(dict(device.status) if device.status else {})

            if device.online:
                dev.last_seen = datetime.utcnow()
            synced += 1

        db.session.commit()
        logger.info("Board %d: synced %d Tuya devices", board.id, synced)
        return synced

    # ── Control ───────────────────────────────────────────────────────────────
    def set_dp(self, board: "Board", device_id: str, dp_code: str, value) -> bool:
        """Send a DP command (e.g. switch=True) to a Tuya device."""
        mgr = build_manager(board)
        if mgr is None:
            return False
        try:
            mgr.send_commands(device_id, [{"code": dp_code, "value": value}])
            self._apply_local(board, device_id, dp_code, value)
            return True
        except Exception as exc:
            logger.error("Tuya set_dp %s.%s=%s failed: %s", device_id, dp_code, value, exc)
            return False

    def toggle_dp(self, board: "Board", device_id: str, dp_code: str) -> bool:
        """Toggle the boolean state of a DP."""
        from app.models.tuya_device import TuyaDevice
        dev = TuyaDevice.query.filter_by(
            board_id=board.id, device_id=device_id
        ).first()
        if not dev:
            return False
        current_val = dev.status.get(dp_code, False)
        return self.set_dp(board, device_id, dp_code, not bool(current_val))

    # ── Private ───────────────────────────────────────────────────────────────
    @staticmethod
    def _apply_local(board: "Board", device_id: str, dp_code: str, value) -> None:
        """Update the cached status immediately so the UI is consistent."""
        from app import db
        from app.models.tuya_device import TuyaDevice
        dev = TuyaDevice.query.filter_by(
            board_id=board.id, device_id=device_id
        ).first()
        if dev:
            s = dev.status
            s[dp_code] = value
            dev.status = s
            try:
                db.session.commit()
            except Exception:
                db.session.rollback()