from __future__ import annotations from enum import StrEnum import re from typing import Any, NamedTuple, Protocol from homeassistant.core import HomeAssistant, State from custom_components.powercalc.common import SourceEntity from custom_components.powercalc.errors import PowercalcSetupError class SubProfileMatcherType(StrEnum): ATTRIBUTE = "attribute" ENTITY_ID = "entity_id" ENTITY_REGISTRY = "entity_registry" ENTITY_STATE = "entity_state" INTEGRATION = "integration" MODEL_ID = "model_id" class SubProfileSelector: def __init__( self, hass: HomeAssistant, config: SubProfileSelectConfig, source_entity: SourceEntity, ) -> None: self._hass = hass self._config = config self._source_entity = source_entity self._matchers: list[SubProfileMatcher] = self._build_matchers() def _build_matchers(self) -> list[SubProfileMatcher]: """Create matchers from json config.""" return [self._create_matcher(matcher_config) for matcher_config in self._config.matchers or []] def select_sub_profile(self, entity_state: State) -> str: """Dynamically tries to select a sub profile depending on the entity state. This method always need to return a sub profile, when nothing is matched it will return a default. """ for matcher in self._matchers: sub_profile = matcher.match(entity_state, self._source_entity) if sub_profile: return sub_profile return self._config.default def get_tracking_entities(self) -> list[str]: """Get additional list of entities to track for state changes.""" return [entity_id for matcher in self._matchers for entity_id in matcher.get_tracking_entities()] def _create_matcher(self, matcher_config: dict[str, Any]) -> SubProfileMatcher: """Create a matcher from json config. Can be extended for more matchers in the future.""" matcher_type: SubProfileMatcherType = matcher_config["type"] match matcher_type: case SubProfileMatcherType.ATTRIBUTE: return AttributeMatcher(matcher_config["attribute"], matcher_config["map"]) case SubProfileMatcherType.ENTITY_STATE: return EntityStateMatcher( self._hass, self._source_entity, matcher_config["entity_id"], matcher_config["map"], ) case SubProfileMatcherType.ENTITY_ID: return EntityIdMatcher(matcher_config["pattern"], matcher_config["profile"]) case SubProfileMatcherType.ENTITY_REGISTRY: return EntityRegistryMatcher( matcher_config["property"], matcher_config["value"], matcher_config["profile"], ) case SubProfileMatcherType.INTEGRATION: return IntegrationMatcher(matcher_config["integration"], matcher_config["profile"]) case SubProfileMatcherType.MODEL_ID: return ModelIdMatcher(matcher_config["model_id"], matcher_config["profile"]) case _: raise PowercalcSetupError(f"Unknown sub profile matcher type: {matcher_type}") class SubProfileSelectConfig(NamedTuple): default: str matchers: list[dict[str, Any]] | None = None class SubProfileMatcher(Protocol): def match(self, entity_state: State, source_entity: SourceEntity) -> str | None: """Returns a sub profile.""" def get_tracking_entities(self) -> list[str]: """Get extra entities to track for state changes.""" return [] class EntityStateMatcher(SubProfileMatcher): def __init__( self, hass: HomeAssistant, source_entity: SourceEntity | None, entity_id: str, mapping: dict[str, str], ) -> None: self._hass = hass if source_entity: entity_id = entity_id.replace( "{{source_object_id}}", source_entity.object_id, ) self._entity_id = entity_id self._mapping = mapping def match(self, entity_state: State, source_entity: SourceEntity) -> str | None: state = self._hass.states.get(self._entity_id) if state is None: return None return self._mapping.get(state.state) def get_tracking_entities(self) -> list[str]: return [self._entity_id] class AttributeMatcher(SubProfileMatcher): def __init__(self, attribute: str, mapping: dict[str, str]) -> None: self._attribute = attribute self._mapping = mapping def match(self, entity_state: State, source_entity: SourceEntity) -> str | None: val = entity_state.attributes.get(self._attribute) if val is None: return None return self._mapping.get(val) class EntityIdMatcher(SubProfileMatcher): def __init__(self, pattern: str, profile: str) -> None: self._pattern = pattern self._profile = profile def match(self, entity_state: State, source_entity: SourceEntity) -> str | None: if re.search(self._pattern, entity_state.entity_id): return self._profile return None class IntegrationMatcher(SubProfileMatcher): def __init__(self, integration: str, profile: str) -> None: self._integration = integration self._profile = profile def match(self, entity_state: State, source_entity: SourceEntity) -> str | None: registry_entry = source_entity.entity_entry if not registry_entry: return None if registry_entry.platform == self._integration: return self._profile return None class EntityRegistryMatcher(SubProfileMatcher): def __init__(self, property_name: str, value: object, profile: str) -> None: self._property_name = property_name self._value = value self._profile = profile def match(self, entity_state: State, source_entity: SourceEntity) -> str | None: registry_entry = source_entity.entity_entry if not registry_entry or not hasattr(registry_entry, self._property_name): return None registry_value = getattr(registry_entry, self._property_name) if registry_value is None: return None if self._matches_registry_value(registry_value): return self._profile return None def _matches_registry_value(self, registry_value: object) -> bool: if registry_value == self._value: return True if isinstance(registry_value, list | set | tuple | frozenset): return self._value in registry_value return str(registry_value) == str(self._value) class ModelIdMatcher(SubProfileMatcher): def __init__(self, model_id: str, profile: str) -> None: self._model_id = model_id self._profile = profile def match(self, entity_state: State, source_entity: SourceEntity) -> str | None: device_entry = source_entity.device_entry if not device_entry: return None if device_entry.model_id == self._model_id: return self._profile return None