HomeAssistantVS/custom_components/powercalc/flow_helper/flows/library.py

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,
        )