HA_Aux_Heater/custom_components/boiler_aux_heater/climate.py

"""Climate entity for Boiler Auxiliary Heater Control."""

from __future__ import annotations

from typing import Any

from homeassistant.components.climate import (
    ClimateEntity,
    ClimateEntityFeature,
    HVACAction,
    HVACMode,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.restore_state import RestoreEntity

from .const import (
    DOMAIN,
    PRESET_AWAY,
    PRESET_BOOST,
    PRESET_ECO,
    PRESET_NORMAL,
    PRESET_SLEEP,
    PRESET_TEMP_MAP,
    PRESETS,
)


async def async_setup_entry(
    hass: HomeAssistant,
    entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up the climate entity."""
    coordinator = hass.data[DOMAIN][entry.entry_id]
    entity = BoilerAuxHeaterClimate(coordinator, entry)
    async_add_entities([entity])


class BoilerAuxHeaterClimate(RestoreEntity, ClimateEntity):
    """Unified climate entity with native presets and dual thermostat logic."""

    _attr_has_entity_name = True
    _attr_translation_key = "boiler_aux_heater"
    _attr_temperature_unit = UnitOfTemperature.CELSIUS
    _attr_hvac_modes = [HVACMode.HEAT, HVACMode.OFF]
    _attr_supported_features = (
        ClimateEntityFeature.TARGET_TEMPERATURE
        | ClimateEntityFeature.PRESET_MODE
    )
    _attr_min_temp = 10
    _attr_max_temp = 70
    _attr_target_temperature_step = 1
    _attr_precision = 0.1
    _enable_turn_on_off_backwards_compat = False

    def __init__(self, coordinator, entry: ConfigEntry) -> None:
        """Initialize the climate entity."""
        self.coordinator = coordinator
        self._attr_unique_id = f"{entry.entry_id}_climate"
        self._attr_preset_modes = list(PRESETS)
        self._attr_preset_mode = PRESET_NORMAL
        self._attr_hvac_mode = HVACMode.HEAT
        self._attr_target_temperature = 50.0
        self._unsub_listener = None
        self._attr_device_info = DeviceInfo(
            identifiers={(DOMAIN, entry.entry_id)},
            name="Boiler Auxiliary Heater",
            manufacturer="Custom",
        )

    async def async_added_to_hass(self) -> None:
        """Restore state and register with coordinator."""
        await super().async_added_to_hass()

        last_state = await self.async_get_last_state()
        if last_state:
            if last_state.state in ("heat", "off"):
                self._attr_hvac_mode = HVACMode(last_state.state)
            attrs = last_state.attributes
            if "temperature" in attrs:
                try:
                    self._attr_target_temperature = float(attrs["temperature"])
                except (ValueError, TypeError):
                    pass
            if "preset_mode" in attrs and attrs["preset_mode"] in PRESETS:
                self._attr_preset_mode = attrs["preset_mode"]

        # Register with coordinator
        self.coordinator.register_climate(self)
        self._unsub_listener = self.coordinator.async_add_listener(
            self._handle_coordinator_update
        )

    async def async_will_remove_from_hass(self) -> None:
        """Clean up listener."""
        if self._unsub_listener:
            self._unsub_listener()

    @callback
    def _handle_coordinator_update(self) -> None:
        """Update entity when coordinator state changes."""
        self.async_write_ha_state()

    @property
    def current_temperature(self) -> float | None:
        """Return current boiler temperature from external sensor."""
        state = self.hass.states.get(self.coordinator.boiler_temp_sensor)
        if state is None or state.state in ("unknown", "unavailable"):
            return None
        try:
            return float(state.state)
        except (ValueError, TypeError):
            return None

    @property
    def hvac_action(self) -> HVACAction:
        """Return current action based on relay state."""
        if self._attr_hvac_mode == HVACMode.OFF:
            return HVACAction.OFF
        if self.coordinator.relay_is_on:
            return HVACAction.HEATING
        return HVACAction.IDLE

    async def async_set_hvac_mode(self, hvac_mode: HVACMode) -> None:
        """Set HVAC mode."""
        self._attr_hvac_mode = hvac_mode
        if hvac_mode == HVACMode.OFF:
            self._attr_preset_mode = None
        self.async_write_ha_state()
        self.coordinator.async_request_refresh()

    async def async_set_temperature(self, **kwargs: Any) -> None:
        """Set target temperature."""
        if ATTR_TEMPERATURE in kwargs:
            self._attr_target_temperature = kwargs[ATTR_TEMPERATURE]
            self.async_write_ha_state()
            self.coordinator.async_request_refresh()

    async def async_set_preset_mode(self, preset_mode: str) -> None:
        """Set preset mode and update target temperature from number entity."""
        self._attr_preset_mode = preset_mode
        self._attr_hvac_mode = HVACMode.HEAT

        # Get target temperature from the corresponding number entity
        temp_key = PRESET_TEMP_MAP.get(preset_mode)
        if temp_key:
            temp = self.coordinator.get_number_value(temp_key)
            if temp is not None:
                self._attr_target_temperature = temp

        # Toggle bath boost based on preset
        if preset_mode == PRESET_BOOST:
            self.coordinator.bath_boost_active = True
        else:
            self.coordinator.bath_boost_active = False

        self.async_write_ha_state()
        self.coordinator.async_request_refresh()