from __future__ import annotations from typing import TYPE_CHECKING, Any from homeassistant.config_entries import ConfigFlowResult from homeassistant.helpers import selector, translation import voluptuous as vol from custom_components.powercalc.const import ( CONF_AVAILABILITY_ENTITY, CONF_FIXED, CONF_MANUFACTURER, CONF_MODE, CONF_MODEL, CONF_POWER, CONF_SELF_USAGE_INCLUDED, CONF_SUB_PROFILE, CONF_VARIABLES, DOMAIN, DUMMY_ENTITY_ID, LIBRARY_URL, CalculationStrategy, ) from custom_components.powercalc.discovery import ( get_power_profile_by_source_device, get_power_profile_by_source_entity, ) from custom_components.powercalc.flow_helper.common import FlowType, PowercalcFormStep, Step from custom_components.powercalc.flow_helper.dynamic_field_builder import build_dynamic_field_schema from custom_components.powercalc.flow_helper.schema import ( SCHEMA_ENERGY_SENSOR_TOGGLE, SCHEMA_UTILITY_METER_TOGGLE, build_sub_profile_schema, ) from custom_components.powercalc.helpers import ( collect_placeholders, iter_related_entity_placeholders, resolve_related_entity_placeholder, ) from custom_components.powercalc.power_profile.library import ModelInfo, ProfileLibrary from custom_components.powercalc.power_profile.power_profile import ( DEVICE_TYPE_DOMAIN, DOMAIN_DEVICE_TYPE_MAPPING, DeviceType, DiscoveryBy, PowerProfile, ) if TYPE_CHECKING: from custom_components.powercalc.config_flow import PowercalcCommonFlow, PowercalcConfigFlow, PowercalcOptionsFlow CONF_CONFIRM_AUTODISCOVERED_MODEL = "confirm_autodisovered_model" SCHEMA_POWER_AUTODISCOVERED = vol.Schema( {vol.Optional(CONF_CONFIRM_AUTODISCOVERED_MODEL, default=True): bool}, ) SCHEMA_POWER_OPTIONS_LIBRARY = vol.Schema( { **SCHEMA_ENERGY_SENSOR_TOGGLE.schema, **SCHEMA_UTILITY_METER_TOGGLE.schema, }, ) SCHEMA_POWER_SMART_SWITCH = vol.Schema( { vol.Optional(CONF_POWER): vol.Coerce(float), vol.Optional(CONF_SELF_USAGE_INCLUDED): selector.BooleanSelector(), }, ) class LibraryFlow: def __init__(self, flow: PowercalcCommonFlow) -> None: self.flow = flow async def async_step_manufacturer( self, user_input: dict[str, Any] | None = None, ) -> ConfigFlowResult: """Ask the user to select the manufacturer.""" async def _create_schema() -> vol.Schema: """Create manufacturer schema.""" library = await ProfileLibrary.factory(self.flow.hass) manufacturers = [ selector.SelectOptionDict(value=manufacturer[0], label=manufacturer[1]) for manufacturer in await library.get_manufacturer_listing( self._get_library_device_types(), self._get_library_discovery_by(), ) ] return vol.Schema( { vol.Required( CONF_MANUFACTURER, default=self.flow.sensor_config.get(CONF_MANUFACTURER), ): selector.SelectSelector( selector.SelectSelectorConfig( options=manufacturers, mode=selector.SelectSelectorMode.DROPDOWN, ), ), }, ) # noinspection PyTypeChecker return await self.flow.handle_form_step( PowercalcFormStep( step=Step.MANUFACTURER, schema=_create_schema, next_step=Step.MODEL, ), user_input, ) async def async_step_model( self, user_input: dict[str, Any] | None = None, ) -> ConfigFlowResult: """Ask the user to select the model.""" def _build_model_label(model_id: str, model_name: str) -> str: if not model_name or model_name == model_id: return model_id return f"{model_id} ({model_name})" async def _validate(user_input: dict[str, Any]) -> dict[str, str]: library = await ProfileLibrary.factory(self.flow.hass) profile = await library.get_profile( ModelInfo( str(self.flow.sensor_config.get(CONF_MANUFACTURER)), str(user_input.get(CONF_MODEL)), ), self.flow.source_entity, process_variables=False, ) self.flow.selected_profile = profile if self.flow.selected_profile and not await self.flow.selected_profile.needs_user_configuration: await self.flow.validate_strategy_config() return user_input async def _create_schema() -> vol.Schema: """Create model schema.""" manufacturer = str(self.flow.sensor_config.get(CONF_MANUFACTURER)) library = await ProfileLibrary.factory(self.flow.hass) models = [ selector.SelectOptionDict(value=model_id, label=_build_model_label(model_id, model_name)) for model_id, model_name in await library.get_model_listing( manufacturer, self._get_library_device_types(), self._get_library_discovery_by(), ) ] model = ( self.flow.selected_profile.model if self.flow.selected_profile else self.flow.sensor_config.get(CONF_MODEL) ) return vol.Schema( { vol.Required( CONF_MODEL, description={"suggested_value": model}, default=model, ): selector.SelectSelector( selector.SelectSelectorConfig( options=models, mode=selector.SelectSelectorMode.DROPDOWN, ), ), }, ) return await self.flow.handle_form_step( PowercalcFormStep( step=Step.MODEL, schema=_create_schema, next_step=Step.POST_LIBRARY, validate_user_input=_validate, form_kwarg={"description_placeholders": {"supported_models_link": LIBRARY_URL}}, ), user_input, ) async def async_step_post_library(self, _: dict[str, Any] | None = None) -> ConfigFlowResult: """ Handles the logic after the user either selected manufacturer/model himself or confirmed autodiscovered. Forwards to the next step in the flow. """ if not self.flow.selected_profile: return self.flow.async_abort(reason="model_not_supported") # pragma: no cover profile_step = await self._async_next_profile_step(self.flow.selected_profile) if profile_step: return profile_step strategy_step = await self._async_next_strategy_step(self.flow.selected_profile) if strategy_step: return strategy_step return await self.flow.flow_handlers[FlowType.GROUP].async_step_assign_groups() # type: ignore[no-any-return] async def _async_next_profile_step(self, profile: PowerProfile) -> ConfigFlowResult | None: """Return the next step needed to complete the profile itself, or None when nothing is left to ask.""" handled_steps = self.flow.handled_steps if Step.LIBRARY_CUSTOM_FIELDS not in handled_steps and profile.has_custom_fields: return await self.async_step_library_custom_fields() if Step.AVAILABILITY_ENTITY not in handled_steps and profile.discovery_by == DiscoveryBy.DEVICE: result = await self.async_step_availability_entity() if result: return result if Step.SUB_PROFILE not in handled_steps and await profile.requires_manual_sub_profile_selection: return await self.async_step_sub_profile() return None async def _async_next_strategy_step(self, profile: PowerProfile) -> ConfigFlowResult | None: """Return the next step needed to configure the calculation strategy, or None when nothing is left to ask.""" handled_steps = self.flow.handled_steps virtual_power_flow = self.flow.flow_handlers[FlowType.VIRTUAL_POWER] if ( Step.SMART_SWITCH not in handled_steps and profile.device_type == DeviceType.SMART_SWITCH and profile.calculation_strategy == CalculationStrategy.FIXED ): return await self.async_step_smart_switch() if Step.FIXED not in handled_steps and profile.needs_fixed_config: # pragma: no cover return await virtual_power_flow.async_step_fixed() # type: ignore[no-any-return] if Step.LINEAR not in handled_steps and profile.needs_linear_config: return await virtual_power_flow.async_step_linear() # type: ignore[no-any-return] if Step.MULTI_SWITCH not in handled_steps and profile.calculation_strategy == CalculationStrategy.MULTI_SWITCH: return await virtual_power_flow.async_step_multi_switch() # type: ignore[no-any-return] return None async def async_step_library_custom_fields(self, user_input: dict[str, Any] | None = None) -> ConfigFlowResult: """Handle the flow for custom fields.""" def _process_user_input(user_input: dict[str, Any]) -> dict[str, Any]: return {CONF_VARIABLES: user_input} form_kwarg: dict[str, Any] | None = None if self.flow.selected_profile and self.flow.selected_profile.documentation_url: form_kwarg = { "description_placeholders": { "documentation_url": self.flow.selected_profile.documentation_url, }, } return await self.flow.handle_form_step( PowercalcFormStep( step=Step.LIBRARY_CUSTOM_FIELDS, schema=build_dynamic_field_schema( self.flow.hass, self.flow.selected_profile, # type: ignore self.flow.source_entity, ), next_step=Step.POST_LIBRARY, validate_user_input=_process_user_input, form_kwarg=form_kwarg, ), user_input, ) async def async_step_sub_profile( self, user_input: dict[str, Any] | None = None, ) -> ConfigFlowResult: """Handle the flow for sub profile selection.""" assert self.flow.selected_profile is not None def _validate(user_input: dict[str, Any]) -> dict[str, str]: return {CONF_MODEL: f"{self.flow.sensor_config.get(CONF_MODEL)}/{user_input.get(CONF_SUB_PROFILE)}"} library = await ProfileLibrary.factory(self.flow.hass) profile = await library.get_profile( ModelInfo( str(self.flow.sensor_config.get(CONF_MANUFACTURER)), str(self.flow.sensor_config.get(CONF_MODEL)), ), self.flow.source_entity, process_variables=False, ) remarks = profile.config_flow_sub_profile_remarks if remarks: remarks = "\n\n" + remarks step = ( Step.SUB_PROFILE_PER_DEVICE if self.flow.selected_profile.discovery_by == DiscoveryBy.DEVICE else Step.SUB_PROFILE ) return await self.flow.handle_form_step( PowercalcFormStep( step=step, schema=await build_sub_profile_schema(profile, self.flow.selected_sub_profile), next_step=Step.POWER_ADVANCED, validate_user_input=_validate, form_kwarg={ "description_placeholders": { "entity_id": self.flow.source_entity_id, "remarks": remarks, }, }, ), user_input, ) async def async_step_sub_profile_per_device( self, user_input: dict[str, Any] | None = None, ) -> ConfigFlowResult: return await self.async_step_sub_profile(user_input) async def async_step_smart_switch(self, user_input: dict[str, Any] | None = None) -> ConfigFlowResult: """Asks the user for the power of connect appliance for the smart switch.""" if self.flow.selected_profile and not self.flow.selected_profile.needs_fixed_config: return self.flow.persist_config_entry() def _validate(user_input: dict[str, Any]) -> dict[str, Any]: return { CONF_SELF_USAGE_INCLUDED: user_input.get(CONF_SELF_USAGE_INCLUDED), CONF_MODE: CalculationStrategy.FIXED, CONF_FIXED: {CONF_POWER: user_input.get(CONF_POWER, 0)}, } self_usage_on = self.flow.selected_profile.standby_power_on if self.flow.selected_profile else 0 return await self.flow.handle_form_step( PowercalcFormStep( step=Step.SMART_SWITCH, schema=SCHEMA_POWER_SMART_SWITCH, validate_user_input=_validate, next_step=Step.POWER_ADVANCED, form_kwarg={"description_placeholders": {"self_usage_power": str(self_usage_on)}}, ), user_input, ) async def async_step_availability_entity(self, user_input: dict[str, Any] | None = None) -> ConfigFlowResult | None: """Handle the flow for availability entity.""" # Auto-resolve availability entity from profile placeholders auto_entity = self._resolve_availability_entity() if auto_entity: self.flow.sensor_config[CONF_AVAILABILITY_ENTITY] = auto_entity self.flow.handled_steps.append(Step.AVAILABILITY_ENTITY) return None domains = DEVICE_TYPE_DOMAIN[self.flow.selected_profile.device_type] # type: ignore entity_selector = self.flow.create_device_entity_selector( list(domains) if isinstance(domains, set) else [domains], ) try: first_entity = entity_selector.config["include_entities"][0] except IndexError: # Skip step if no entities are available self.flow.handled_steps.append(Step.AVAILABILITY_ENTITY) return None return await self.flow.handle_form_step( PowercalcFormStep( step=Step.AVAILABILITY_ENTITY, schema=vol.Schema( { vol.Optional(CONF_AVAILABILITY_ENTITY, default=first_entity): entity_selector, }, ), next_step=Step.POST_LIBRARY, ), user_input, ) def _resolve_availability_entity(self) -> str | None: """Try to auto-resolve an availability entity from profile placeholders.""" profile = self.flow.selected_profile device_entry = self.flow.source_entity.device_entry if self.flow.source_entity else None if not profile or not device_entry: return None for placeholder in iter_related_entity_placeholders(collect_placeholders(profile.json_data)): entity = resolve_related_entity_placeholder( self.flow.hass, placeholder, source_entity=self.flow.source_entity, ) if entity: return entity return None def _get_library_device_types(self) -> set[DeviceType] | None: """Determine which device types should be shown in the library selectors.""" if self._get_library_discovery_by() == DiscoveryBy.DEVICE: return None if self.flow.source_entity: return DOMAIN_DEVICE_TYPE_MAPPING.get(self.flow.source_entity.domain, set()) return None # pragma: no cover def _get_library_discovery_by(self) -> DiscoveryBy | None: """Determine whether listing should be filtered by discovery mode.""" if self.flow.source_entity and self.flow.source_entity.entity_id == DUMMY_ENTITY_ID: return DiscoveryBy.DEVICE return None class LibraryConfigFlow(LibraryFlow): def __init__(self, flow: PowercalcConfigFlow) -> None: super().__init__(flow) self.flow: PowercalcConfigFlow = flow async def async_step_library_multi_profile( self, user_input: dict[str, Any] | None = None, ) -> ConfigFlowResult: """This step gets executed when multiple profiles are found for the source entity.""" if user_input is not None: selected_model: str = user_input.get(CONF_MODEL) # type: ignore selected_profile = self.flow.discovered_profiles.get(selected_model) if selected_profile is None: # pragma: no cover return self.flow.async_abort(reason="invalid_profile") self.flow.selected_profile = selected_profile self.flow.sensor_config.update( { CONF_MANUFACTURER: selected_profile.manufacturer, CONF_MODEL: selected_profile.model, }, ) return await self.async_step_post_library(user_input) schema = vol.Schema( { vol.Required(CONF_MODEL): selector.SelectSelector( selector.SelectSelectorConfig( options=[ selector.SelectOptionDict( value=profile.unique_id, label=profile.model, ) for profile in self.flow.discovered_profiles.values() ], mode=selector.SelectSelectorMode.DROPDOWN, ), ), }, ) manufacturer = str(self.flow.sensor_config.get(CONF_MANUFACTURER)) model = str(self.flow.sensor_config.get(CONF_MODEL)) return self.flow.async_show_form( step_id=Step.LIBRARY_MULTI_PROFILE, data_schema=schema, description_placeholders={ "library_link": f"{LIBRARY_URL}/?manufacturer={manufacturer}", "manufacturer": manufacturer, "model": model, }, last_step=False, ) async def async_step_library( self, user_input: dict[str, Any] | None = None, ) -> ConfigFlowResult: """Try to autodiscover manufacturer/model first. Ask the user to confirm this or forward to manual library selection. """ if user_input is not None: return await self._handle_library_confirmation(user_input) await self._async_autodiscover_profile() if not self.flow.selected_profile: return await self.async_step_manufacturer() return self._show_autodiscovered_profile_form() async def _handle_library_confirmation(self, user_input: dict[str, Any]) -> ConfigFlowResult: """Handle the user's response to an autodiscovered library profile.""" if not user_input.get(CONF_CONFIRM_AUTODISCOVERED_MODEL) or not self.flow.selected_profile: return await self.async_step_manufacturer() self.flow.sensor_config.update( { CONF_MANUFACTURER: self.flow.selected_profile.manufacturer, CONF_MODEL: self.flow.selected_profile.model, }, ) return await self.async_step_post_library(user_input) async def _async_autodiscover_profile(self) -> None: """Populate the selected profile from the source entity when possible.""" if ( not self.flow.source_entity or not self.flow.source_entity.entity_entry or self.flow.selected_profile is not None ): return self.flow.selected_profile = await get_power_profile_by_source_entity(self.flow.hass, self.flow.source_entity) if self.flow.selected_profile is None and self.flow.source_entity.device_entry: self.flow.selected_profile = await get_power_profile_by_source_device( self.flow.hass, self.flow.source_entity, ) def _show_autodiscovered_profile_form(self) -> ConfigFlowResult: """Show the confirmation form for an autodiscovered library profile.""" profile = self.flow.selected_profile assert profile is not None return self.flow.async_show_form( step_id=Step.LIBRARY, description_placeholders=self._build_library_description_placeholders(profile), data_schema=SCHEMA_POWER_AUTODISCOVERED, errors={}, last_step=False, ) def _build_library_description_placeholders(self, profile: PowerProfile) -> dict[str, Any]: """Build the placeholders for the autodiscovered profile confirmation form.""" return { "remarks": self._build_library_remarks(profile), "manufacturer": profile.manufacturer, "model": profile.model, "source": self._get_profile_source(profile), } def _build_library_remarks(self, profile: PowerProfile) -> str | None: """Build the remarks text for the autodiscovered profile confirmation form.""" remarks = self._get_conditional_remarks() if remarks: remarks = "\n\n" + remarks if profile.documentation_url: return (remarks or "") + f"\n\n[Documentation]({profile.documentation_url})" return remarks def _get_profile_source(self, profile: PowerProfile) -> str: """Build the autodiscovery source description.""" translations = translation.async_get_cached_translations( self.flow.hass, self.flow.hass.config.language, "common", DOMAIN, ) if ( profile.discovery_by == DiscoveryBy.DEVICE and self.flow.source_entity and self.flow.source_entity.device_entry ): label = translations.get(f"component.{DOMAIN}.common.source_device") return f"{label}: {self.flow.source_entity.device_entry.name}" return f"{translations.get(f'component.{DOMAIN}.common.source_entity')}: {self.flow.source_entity_id}" def _get_conditional_remarks(self) -> str | None: """Get discovery remarks, only showing them if required entities are missing.""" profile = self.flow.selected_profile if not profile: return None # pragma: no cover remarks = profile.config_flow_discovery_remarks if not remarks: return None # Check if all entity_by_* placeholders can be resolved from the device device_entry = self.flow.source_entity.device_entry if self.flow.source_entity else None if not device_entry: return remarks related_placeholders = iter_related_entity_placeholders(collect_placeholders(profile.json_data)) all_resolved = all( resolve_related_entity_placeholder( self.flow.hass, placeholder, source_entity=self.flow.source_entity, ) for placeholder in related_placeholders ) return None if all_resolved else remarks class LibraryOptionsFlow(LibraryFlow): def __init__(self, flow: PowercalcOptionsFlow) -> None: super().__init__(flow) self.flow: PowercalcOptionsFlow = flow async def async_step_library_options(self, user_input: dict[str, Any] | None = None) -> ConfigFlowResult: """Handle the basic options flow.""" self.flow.is_library_flow = True self.flow.selected_sub_profile = self.flow.selected_profile.sub_profile # type: ignore if user_input is not None: return await self.async_step_manufacturer() return self.flow.async_show_form( step_id=Step.LIBRARY_OPTIONS, description_placeholders={ "manufacturer": self.flow.selected_profile.manufacturer, # type: ignore "model": self.flow.selected_profile.model, # type: ignore }, last_step=False, )