sillyhome-next-dev/app/ha/discovery.py

from __future__ import annotations

from enum import StrEnum

from pydantic import BaseModel

from app.ha.models import HaEntitySummary


class EntityRole(StrEnum):
    MEASUREMENT = "measurement"
    BINARY_CONTEXT = "binary_context"
    CONTEXT = "context"
    ACTUATOR = "actuator"
    UNSUPPORTED = "unsupported"


class DiscoveredEntity(BaseModel):
    entity_id: str
    domain: str
    device_class: str | None = None
    state_class: str | None = None
    unit_of_measurement: str | None = None
    category: str
    role: EntityRole
    learnable: bool
    reason: str


_MEASUREMENT_CLASSES = frozenset({
    "apparent_power",
    "atmospheric_pressure",
    "battery",
    "carbon_dioxide",
    "carbon_monoxide",
    "current",
    "distance",
    "duration",
    "energy",
    "frequency",
    "gas",
    "humidity",
    "illuminance",
    "moisture",
    "monetary",
    "nitrogen_dioxide",
    "nitrogen_monoxide",
    "nitrous_oxide",
    "ozone",
    "pm1",
    "pm10",
    "pm25",
    "power",
    "precipitation",
    "pressure",
    "reactive_power",
    "signal_strength",
    "sound_pressure",
    "speed",
    "sulphur_dioxide",
    "temperature",
    "volatile_organic_compounds",
    "voltage",
    "volume",
    "volume_flow_rate",
    "water",
    "weight",
    "wind_speed",
})
_BINARY_CONTEXT_CLASSES = frozenset({
    "door",
    "garage_door",
    "lock",
    "motion",
    "occupancy",
    "opening",
    "presence",
    "problem",
    "safety",
    "smoke",
    "sound",
    "vibration",
    "window",
})
_ACTUATOR_DOMAINS = frozenset({
    "button",
    "climate",
    "cover",
    "fan",
    "humidifier",
    "input_boolean",
    "input_button",
    "lock",
    "light",
    "number",
    "siren",
    "switch",
    "valve",
})
_CONTEXT_DOMAINS = frozenset({
    "device_tracker",
    "input_boolean",
    "input_datetime",
    "input_number",
    "input_select",
    "person",
    "sun",
    "weather",
    "zone",
})
_LEARNABLE_CONTEXT_DOMAINS = frozenset({
    "device_tracker",
    "input_boolean",
    "input_number",
    "input_select",
    "person",
    "weather",
})
_NUMERIC_STATE_CLASSES = frozenset({"measurement", "total", "total_increasing"})


def classify_entity(entity: HaEntitySummary) -> DiscoveredEntity:
    if entity.domain == "sensor" and (
        entity.state_class in _NUMERIC_STATE_CLASSES
        or entity.device_class in _MEASUREMENT_CLASSES
        or entity.unit_of_measurement is not None
    ):
        return _result(
            entity,
            EntityRole.MEASUREMENT,
            category=_measurement_category(entity),
            learnable=True,
            reason="Numerischer Messsensor für Zeitreihen und Training.",
        )

    if entity.domain == "binary_sensor" and entity.device_class in _BINARY_CONTEXT_CLASSES:
        return _result(
            entity,
            EntityRole.BINARY_CONTEXT,
            category=_binary_category(entity),
            learnable=True,
            reason="Binärer Kontextsensor für Zustands- und Anwesenheitsmuster.",
        )

    if entity.domain in _CONTEXT_DOMAINS:
        learnable = entity.domain in _LEARNABLE_CONTEXT_DOMAINS
        return _result(
            entity,
            EntityRole.CONTEXT,
            category=_context_category(entity),
            learnable=learnable,
            reason=(
                "Kontextquelle für Training und Erklärungen."
                if learnable
                else "Kontextquelle ohne direkte Trainingsfreigabe."
            ),
        )

    if entity.domain in _ACTUATOR_DOMAINS:
        return _result(
            entity,
            EntityRole.ACTUATOR,
            category=_actuator_category(entity),
            learnable=False,
            reason="Aktor ist ein mögliches Automationsziel, aber kein Trainingssensor.",
        )

    return _result(
        entity,
        EntityRole.UNSUPPORTED,
        category="unsupported",
        learnable=False,
        reason="Entity-Typ ist noch nicht für Lernen oder Automationen klassifiziert.",
    )


def discover_entities(
    entities: list[HaEntitySummary],
    domains: set[str] | None = None,
    learnable: bool | None = None,
) -> list[DiscoveredEntity]:
    normalized_domains = {domain.strip().lower() for domain in domains or set() if domain.strip()}
    discovered = [classify_entity(entity) for entity in entities]
    return [
        entity
        for entity in discovered
        if (not normalized_domains or entity.domain in normalized_domains)
        and (learnable is None or entity.learnable is learnable)
    ]


def _result(
    entity: HaEntitySummary,
    role: EntityRole,
    *,
    category: str,
    learnable: bool,
    reason: str,
) -> DiscoveredEntity:
    return DiscoveredEntity(
        entity_id=entity.entity_id,
        domain=entity.domain,
        device_class=entity.device_class,
        state_class=entity.state_class,
        unit_of_measurement=entity.unit_of_measurement,
        category=category,
        role=role,
        learnable=learnable,
        reason=reason,
    )


def _actuator_category(entity: HaEntitySummary) -> str:
    if entity.domain == "light":
        return "light"
    if entity.domain == "switch":
        return "switch_socket"
    if entity.domain == "button" or entity.domain == "input_button":
        return "button"
    if entity.domain == "cover":
        return "cover_shutter"
    if entity.domain == "climate":
        return "heating"
    if entity.domain == "lock":
        return "lock"
    if entity.domain == "fan":
        return "fan"
    if entity.domain in {"input_boolean", "number"}:
        return "helper"
    return entity.domain


def _measurement_category(entity: HaEntitySummary) -> str:
    device_class = entity.device_class or ""
    if device_class == "illuminance":
        return "brightness"
    if device_class == "temperature":
        return "temperature"
    if device_class in {"humidity", "moisture"}:
        return "humidity"
    if device_class in {"power", "energy", "current", "voltage"}:
        return "energy_power"
    if device_class in {"battery", "signal_strength"}:
        return "diagnostic"
    return "measurement"


def _binary_category(entity: HaEntitySummary) -> str:
    device_class = entity.device_class or ""
    if device_class in {"motion", "occupancy", "presence"}:
        return "presence_motion"
    if device_class in {"door", "garage_door", "opening", "window"}:
        return "opening"
    if device_class in {"smoke", "safety", "problem"}:
        return "safety"
    return "binary"


def _context_category(entity: HaEntitySummary) -> str:
    if entity.domain.startswith("input_"):
        return "helper"
    if entity.domain in {"person", "device_tracker", "zone"}:
        return "presence_location"
    return entity.domain