home-automation/app/services/ha_discovery.py

"""HA MQTT Discovery service — build and publish discovery configs + state.

This module handles:

1. Building HA MQTT Discovery payloads for each ``ExposableEntity``.
2. Publishing discovery configs (retained) for enabled entities.
3. Clearing (empty-payload) discovery configs for disabled entities.
4. Publishing current entity state values.
5. A per-device helper for use in the Modbus poll loop.

Design
------
- **No-op when MQTT / discovery is not configured**: every public function
  checks ``ha_discovery_enabled`` and ``mqtt_manager.is_connected`` before
  doing any work.  Callers never need to guard these conditions.
- **Stable unique_id**: derived from device ``uuid`` + metric ``key``, never
  from mutable fields like ``friendly_name`` or auto-increment DB ids.
- **Discovery topic format**: ``<discovery_prefix>/<component>/<node_id>/<object_id>/config``
  where ``node_id`` is the device uuid (slugified to be safe) and
  ``object_id`` is a uuid-prefixed stable string.  Uses ``ha_discovery_prefix``
  (default ``"homeassistant"``), which must stay under the HA discovery namespace.
- **State topic format**: ``<state_prefix>/<component>/<node_id>/<object_id>/state``
  Uses ``ha_state_topic_prefix`` (default ``"home_automation"``), separate from
  the discovery namespace so HA discovery and state/availability topics live under
  different prefixes.
- **Availability topic**: ``<state_prefix>/modbus/<node_id>/availability`` (shared
  across all entities of the same device; publishes "online"/"offline").  Also uses
  the state prefix, not the discovery prefix.
- **best-effort**: functions catch all exceptions internally so that callers
  (e.g. the Modbus poll loop) never crash due to MQTT publish errors.
"""

from __future__ import annotations

import json
import logging
from typing import Any

from sqlalchemy.orm import Session

from app.integrations.expose import ExposableEntity, build_catalog
from app.integrations.mqtt import mqtt_manager
from app.services.config_page import build_runtime_settings

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _should_publish(settings: Any) -> bool:
    """Return True only if MQTT and HA Discovery are both enabled and connected."""
    return bool(
        settings.mqtt_enabled
        and settings.ha_discovery_enabled
        and mqtt_manager.is_connected
    )


def _node_id(device_uuid: str) -> str:
    """Stable MQTT node_id derived from device uuid (replace hyphens for safety)."""
    return device_uuid.replace("-", "_")


def _object_id(entity: ExposableEntity) -> str:
    """Stable MQTT object_id derived from entity key (dots + hyphens → underscores)."""
    return entity.key.replace(".", "_").replace("-", "_")


def _discovery_topic(entity: ExposableEntity, prefix: str) -> str:
    """Build the HA Discovery config topic for *entity*.

    Format: ``<prefix>/<component>/<node_id>/<object_id>/config``
    """
    node = _node_id(entity.device.identifiers[1])  # device uuid
    obj = _object_id(entity)
    return f"{prefix}/{entity.component}/{node}/{obj}/config"


def _state_topic(entity: ExposableEntity, prefix: str) -> str:
    """Build the state publish topic for *entity*.

    Format: ``<prefix>/<component>/<node_id>/<object_id>/state``
    """
    node = _node_id(entity.device.identifiers[1])
    obj = _object_id(entity)
    return f"{prefix}/{entity.component}/{node}/{obj}/state"


def _availability_topic(device_uuid: str, prefix: str) -> str:
    """Shared availability topic for all entities of a device.

    Format: ``<prefix>/modbus/<node_id>/availability``
    """
    node = _node_id(device_uuid)
    return f"{prefix}/modbus/{node}/availability"


def _unique_id(entity: ExposableEntity) -> str:
    """Stable unique_id — device uuid + metric key (never from mutable fields)."""
    device_uuid = entity.device.identifiers[1]
    # entity.key is already "modbus.<uuid>.<metric_key>" — use it as the seed
    return f"{device_uuid}_{entity.key.replace('.', '_')}"


# ---------------------------------------------------------------------------
# Public: build discovery payload
# ---------------------------------------------------------------------------


def build_discovery_payload(
    entity: ExposableEntity,
    discovery_prefix: str,
    state_prefix: str | None = None,
) -> tuple[str, dict[str, Any]]:
    """Build the HA MQTT Discovery config topic and payload dict for *entity*.

    Parameters
    ----------
    entity:
        An ``ExposableEntity`` (from ``build_catalog``).
    discovery_prefix:
        The HA Discovery config topic prefix (e.g. ``"homeassistant"``).
        This determines where HA looks for the discovery config topic.
    state_prefix:
        The prefix used for state and availability topics (e.g.
        ``"home_automation"``).  When omitted or ``None``, falls back to
        ``discovery_prefix`` for backward compatibility (e.g. in tests that
        pass only one prefix).

    Returns
    -------
    tuple[str, dict]
        ``(topic, config_dict)`` — the config topic and the HA Discovery
        payload (not yet JSON-serialised).
    """
    if state_prefix is None:
        state_prefix = discovery_prefix

    device_uuid = entity.device.identifiers[1]
    avail_topic = _availability_topic(device_uuid, state_prefix)
    state_t = _state_topic(entity, state_prefix)
    topic = _discovery_topic(entity, discovery_prefix)

    config: dict[str, Any] = {
        "unique_id": _unique_id(entity),
        "name": entity.name,
        "state_topic": state_t,
        "device": {
            "identifiers": list(entity.device.identifiers),
            "name": entity.device.name,
        },
    }

    # Only declare an availability topic for devices that actually publish an
    # online/offline heartbeat (e.g. Modbus, via its "online" binary_sensor).
    # Devices without a heartbeat (e.g. energy-cost) omit availability so HA
    # treats their entities as always-available — otherwise HA would mark them
    # ``unavailable`` forever even though their state is being published.
    if entity.device.provides_availability:
        config["availability"] = [{"topic": avail_topic}]
        config["availability_mode"] = "all"

    # Component-specific fields
    if entity.component == "binary_sensor":
        # "online" binary sensor: payload ON/OFF
        config["payload_on"] = "ON"
        config["payload_off"] = "OFF"
        if entity.device_class:
            config["device_class"] = entity.device_class
    else:
        # sensor (and others)
        if entity.device_class:
            config["device_class"] = entity.device_class
        if entity.unit:
            config["unit_of_measurement"] = entity.unit
        if entity.state_class:
            config["state_class"] = entity.state_class

    return topic, config


# ---------------------------------------------------------------------------
# Public: publish discovery
# ---------------------------------------------------------------------------


def publish_discovery(session: Session) -> None:
    """Publish HA Discovery configs for all entities in the catalog.

    - Enabled entities: publish retained JSON config payload.
    - Disabled entities: publish empty (b"") payload to clear any previously
      retained config from HA.

    No-op if MQTT / discovery is not enabled or the client is not connected.
    All MQTT errors are caught internally — this function never raises.
    """
    from app.config import get_settings
    settings = build_runtime_settings(session, get_settings())
    if not _should_publish(settings):
        logger.debug("publish_discovery: skipped (MQTT/Discovery not enabled or not connected)")
        return

    discovery_prefix = settings.ha_discovery_prefix
    state_prefix = settings.ha_state_topic_prefix

    try:
        catalog = build_catalog(session)
    except Exception:
        logger.exception("publish_discovery: failed to build catalog; aborting")
        return

    for entry in catalog:
        entity = entry.entity
        try:
            topic, config = build_discovery_payload(entity, discovery_prefix, state_prefix)
            if entry.enabled:
                payload = json.dumps(config)
                mqtt_manager.publish(topic, payload, retain=True)
                logger.debug(
                    "publish_discovery: published config for %r → %s", entity.key, topic
                )
            else:
                # Clear retained config for disabled entities.
                mqtt_manager.publish(topic, b"", retain=True)
                logger.debug(
                    "publish_discovery: cleared config for disabled entity %r → %s",
                    entity.key,
                    topic,
                )
        except Exception:
            logger.exception(
                "publish_discovery: error processing entity %r; continuing", entity.key
            )


# ---------------------------------------------------------------------------
# Public: publish states
# ---------------------------------------------------------------------------


def publish_states(session: Session) -> None:
    """Publish current state values for all enabled entities.

    Calls each entity's ``value_getter`` to obtain the current value and
    publishes it to the entity's state topic.  Also publishes availability
    (online/offline) for each device.

    No-op if MQTT / discovery is not enabled or the client is not connected.
    All errors are caught internally — this function never raises.
    """
    from app.config import get_settings
    settings = build_runtime_settings(session, get_settings())
    if not _should_publish(settings):
        logger.debug("publish_states: skipped (MQTT/Discovery not enabled or not connected)")
        return

    state_prefix = settings.ha_state_topic_prefix

    try:
        catalog = build_catalog(session)
    except Exception:
        logger.exception("publish_states: failed to build catalog; aborting")
        return

    # Collect enabled entries and publish
    for entry in catalog:
        if not entry.enabled:
            continue
        entity = entry.entity
        try:
            _publish_entity_state(entity, state_prefix, session)
        except Exception:
            logger.exception(
                "publish_states: error publishing state for %r; continuing", entity.key
            )


def _publish_entity_state(
    entity: ExposableEntity,
    prefix: str,
    session: Session,
) -> None:
    """Publish one entity's current state value to its state topic.

    Also publishes the availability topic for ``binary_sensor`` "online" entities.
    """
    state_t = _state_topic(entity, prefix)
    device_uuid = entity.device.identifiers[1]

    if entity.component == "binary_sensor" and "online" in entity.key:
        # The online sensor represents device availability.
        # Ask the value_getter for the current ON/OFF string, then also publish
        # the shared availability topic (online/offline).
        raw_value = None
        if entity.value_getter is not None:
            try:
                raw_value = entity.value_getter(session)
            except Exception:
                logger.exception("value_getter raised for online entity %r", entity.key)
        # Default to offline when no reading is available.
        online = (raw_value == "ON")
        avail_payload = "online" if online else "offline"
        avail_topic = _availability_topic(device_uuid, prefix)
        mqtt_manager.publish(avail_topic, avail_payload, retain=False)
        # The state of the binary_sensor itself
        state_payload = "ON" if online else "OFF"
        mqtt_manager.publish(state_t, state_payload, retain=False)
    else:
        # Regular sensor: call value_getter(session) if available.
        value = None
        if entity.value_getter is not None:
            try:
                value = entity.value_getter(session)
            except Exception:
                logger.exception("value_getter raised for entity %r", entity.key)

        if value is None:
            logger.debug("publish_states: no value for %r — skipping state publish", entity.key)
            return

        mqtt_manager.publish(state_t, str(value), retain=False)


# ---------------------------------------------------------------------------
# Public: per-device state push (called by modbus_poll after each poll)
# ---------------------------------------------------------------------------


def publish_device_state(session: Session, device: Any) -> None:
    """Publish state + availability for all enabled entities of *device*.

    Called by ``modbus_poll.poll_device`` after a successful (or failed) poll.
    Publishes:
    - The current availability topic ("online" / "offline") for the device.
    - The state topic for each **enabled** entity of the device.

    No-op if MQTT / discovery is not enabled or the client is not connected.
    All errors are caught internally — never raises.
    """
    from app.config import get_settings
    settings = build_runtime_settings(session, get_settings())
    if not _should_publish(settings):
        return

    state_prefix = settings.ha_state_topic_prefix
    device_uuid = device.uuid
    online = bool(device.last_poll_ok)

    # Publish availability topic first.
    try:
        avail_topic = _availability_topic(device_uuid, state_prefix)
        avail_payload = "online" if online else "offline"
        mqtt_manager.publish(avail_topic, avail_payload, retain=False)
    except Exception:
        logger.exception(
            "publish_device_state: failed to publish availability for device uuid=%s", device_uuid
        )

    if not online:
        # Device is offline — no point pushing stale state values.
        return

    # Publish state for each enabled entity of this device.
    try:
        catalog = build_catalog(session)
    except Exception:
        logger.exception(
            "publish_device_state: failed to build catalog for device uuid=%s", device_uuid
        )
        return

    for entry in catalog:
        if not entry.enabled:
            continue
        entity = entry.entity
        # Only process entities belonging to this device.
        if entity.device.identifiers[1] != device_uuid:
            continue
        # Skip the online binary_sensor itself (availability handled above).
        if entity.component == "binary_sensor" and "online" in entity.key:
            # Publish the binary_sensor state too.
            try:
                state_t = _state_topic(entity, state_prefix)
                mqtt_manager.publish(state_t, "ON", retain=False)
            except Exception:
                logger.exception(
                    "publish_device_state: error publishing online sensor state for %r",
                    entity.key,
                )
            continue

        # Regular sensor — call value_getter(session).
        try:
            value = None
            if entity.value_getter is not None:
                value = entity.value_getter(session)
            if value is None:
                continue
            state_t = _state_topic(entity, state_prefix)
            mqtt_manager.publish(state_t, str(value), retain=False)
        except Exception:
            logger.exception(
                "publish_device_state: error publishing state for entity %r", entity.key
            )


def clear_device_discovery(session: Session, device_uuid: str) -> None:
    """Best-effort: send empty retained payloads to all HA Discovery config topics
    for the given device, effectively removing the device's entities from HA.

    This must be called **before** the device rows are deleted from the DB, so
    that ``build_catalog`` can still enumerate the device's entities.

    Behaviour
    ---------
    - No-op if MQTT / HA Discovery is not enabled or the MQTT client is not
      connected.
    - All exceptions are caught internally; this function never raises.
      The caller proceeds with the DB deletion regardless of MQTT outcome.

    Parameters
    ----------
    session:
        Active SQLAlchemy session (device must still exist in DB at call time).
    device_uuid:
        UUID string of the device being deleted.
    """
    from app.config import get_settings
    settings = build_runtime_settings(session, get_settings())
    if not _should_publish(settings):
        logger.debug(
            "clear_device_discovery: skipped (MQTT/Discovery not enabled or not connected)"
        )
        return

    discovery_prefix = settings.ha_discovery_prefix
    state_prefix = settings.ha_state_topic_prefix

    try:
        catalog = build_catalog(session)
    except Exception:
        logger.exception(
            "clear_device_discovery: failed to build catalog for device uuid=%s; skipping HA cleanup",
            device_uuid,
        )
        return

    cleared = 0
    for entry in catalog:
        entity = entry.entity
        # Only clear entities belonging to this device.
        if entity.device.identifiers[1] != device_uuid:
            continue
        try:
            topic, _config = build_discovery_payload(entity, discovery_prefix, state_prefix)
            mqtt_manager.publish(topic, b"", retain=True)
            cleared += 1
            logger.debug(
                "clear_device_discovery: cleared config topic for entity %r → %s",
                entity.key,
                topic,
            )
        except Exception:
            logger.exception(
                "clear_device_discovery: error clearing entity %r; continuing", entity.key
            )

    logger.info(
        "clear_device_discovery: cleared %d HA discovery topic(s) for device uuid=%s",
        cleared,
        device_uuid,
    )


def publish_device_offline(session: Session, device_uuid: str) -> None:
    """Publish an "offline" availability payload for *device_uuid*.

    Called by ``modbus_poll.poll_device`` when a poll fails, to immediately
    reflect the device going offline in HA (before the next full state sweep).

    No-op if MQTT / discovery is not enabled or the client is not connected.
    Never raises.
    """
    from app.config import get_settings
    settings = build_runtime_settings(session, get_settings())
    if not _should_publish(settings):
        return

    try:
        state_prefix = settings.ha_state_topic_prefix
        avail_topic = _availability_topic(device_uuid, state_prefix)
        mqtt_manager.publish(avail_topic, "offline", retain=False)
    except Exception:
        logger.exception(
            "publish_device_offline: failed for device uuid=%s", device_uuid
        )