core/homeassistant/components/rainmachine/switch.py

"""This component provides support for RainMachine programs and zones."""
from __future__ import annotations

from collections.abc import Coroutine
from datetime import datetime
from typing import Any

from regenmaschine.controller import Controller
from regenmaschine.errors import RequestError
import voluptuous as vol

from homeassistant.components.switch import SwitchEntity
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_ID
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import config_validation as cv, entity_platform
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator

from . import RainMachineEntity, async_update_programs_and_zones
from .const import (
    CONF_ZONE_RUN_TIME,
    DATA_CONTROLLER,
    DATA_COORDINATOR,
    DATA_PROGRAMS,
    DATA_ZONES,
    DEFAULT_ZONE_RUN,
    DOMAIN,
    LOGGER,
)

ATTR_AREA = "area"
ATTR_CS_ON = "cs_on"
ATTR_CURRENT_CYCLE = "current_cycle"
ATTR_CYCLES = "cycles"
ATTR_DELAY = "delay"
ATTR_DELAY_ON = "delay_on"
ATTR_FIELD_CAPACITY = "field_capacity"
ATTR_NEXT_RUN = "next_run"
ATTR_NO_CYCLES = "number_of_cycles"
ATTR_PRECIP_RATE = "sprinkler_head_precipitation_rate"
ATTR_RESTRICTIONS = "restrictions"
ATTR_SLOPE = "slope"
ATTR_SOAK = "soak"
ATTR_SOIL_TYPE = "soil_type"
ATTR_SPRINKLER_TYPE = "sprinkler_head_type"
ATTR_STATUS = "status"
ATTR_SUN_EXPOSURE = "sun_exposure"
ATTR_TIME_REMAINING = "time_remaining"
ATTR_VEGETATION_TYPE = "vegetation_type"
ATTR_ZONES = "zones"

CONF_PROGRAM_ID = "program_id"
CONF_SECONDS = "seconds"
CONF_ZONE_ID = "zone_id"

DEFAULT_ICON = "mdi:water"

DAYS = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"]

RUN_STATUS_MAP = {0: "Not Running", 1: "Running", 2: "Queued"}

SOIL_TYPE_MAP = {
    0: "Not Set",
    1: "Clay Loam",
    2: "Silty Clay",
    3: "Clay",
    4: "Loam",
    5: "Sandy Loam",
    6: "Loamy Sand",
    7: "Sand",
    8: "Sandy Clay",
    9: "Silt Loam",
    10: "Silt",
    99: "Other",
}

SLOPE_TYPE_MAP = {
    0: "Not Set",
    1: "Flat",
    2: "Moderate",
    3: "High",
    4: "Very High",
    99: "Other",
}

SPRINKLER_TYPE_MAP = {
    0: "Not Set",
    1: "Popup Spray",
    2: "Rotors Low Rate",
    3: "Surface Drip",
    4: "Bubblers Drip",
    5: "Rotors High Rate",
    99: "Other",
}

SUN_EXPOSURE_MAP = {0: "Not Set", 1: "Full Sun", 2: "Partial Shade", 3: "Full Shade"}

VEGETATION_MAP = {
    0: "Not Set",
    1: "Not Set",
    2: "Cool Season Grass",
    3: "Fruit Trees",
    4: "Flowers",
    5: "Vegetables",
    6: "Citrus",
    7: "Bushes",
    9: "Drought Tolerant Plants",
    10: "Warm Season Grass",
    11: "Trees",
    99: "Other",
}

SWITCH_TYPE_PROGRAM = "program"
SWITCH_TYPE_ZONE = "zone"


async def async_setup_entry(
    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
) -> None:
    """Set up RainMachine switches based on a config entry."""
    platform = entity_platform.async_get_current_platform()

    alter_program_schema = {vol.Required(CONF_PROGRAM_ID): cv.positive_int}
    alter_zone_schema = {vol.Required(CONF_ZONE_ID): cv.positive_int}

    for service_name, schema, method in (
        ("disable_program", alter_program_schema, "async_disable_program"),
        ("disable_zone", alter_zone_schema, "async_disable_zone"),
        ("enable_program", alter_program_schema, "async_enable_program"),
        ("enable_zone", alter_zone_schema, "async_enable_zone"),
        (
            "pause_watering",
            {vol.Required(CONF_SECONDS): cv.positive_int},
            "async_pause_watering",
        ),
        (
            "start_program",
            {vol.Required(CONF_PROGRAM_ID): cv.positive_int},
            "async_start_program",
        ),
        (
            "start_zone",
            {
                vol.Required(CONF_ZONE_ID): cv.positive_int,
                vol.Optional(
                    CONF_ZONE_RUN_TIME, default=DEFAULT_ZONE_RUN
                ): cv.positive_int,
            },
            "async_start_zone",
        ),
        ("stop_all", {}, "async_stop_all"),
        (
            "stop_program",
            {vol.Required(CONF_PROGRAM_ID): cv.positive_int},
            "async_stop_program",
        ),
        ("stop_zone", {vol.Required(CONF_ZONE_ID): cv.positive_int}, "async_stop_zone"),
        ("unpause_watering", {}, "async_unpause_watering"),
    ):
        platform.async_register_entity_service(service_name, schema, method)

    controller = hass.data[DOMAIN][DATA_CONTROLLER][entry.entry_id]
    programs_coordinator = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][
        DATA_PROGRAMS
    ]
    zones_coordinator = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id][DATA_ZONES]

    entities: list[RainMachineProgram | RainMachineZone] = []

    for uid, program in programs_coordinator.data.items():
        entities.append(
            RainMachineProgram(
                programs_coordinator, controller, uid, program["name"], entry
            )
        )
    for uid, zone in zones_coordinator.data.items():
        entities.append(
            RainMachineZone(zones_coordinator, controller, uid, zone["name"], entry)
        )

    async_add_entities(entities)


class RainMachineSwitch(RainMachineEntity, SwitchEntity):
    """A class to represent a generic RainMachine switch."""

    _attr_icon = DEFAULT_ICON

    def __init__(
        self,
        coordinator: DataUpdateCoordinator,
        controller: Controller,
        uid: int,
        name: str,
        entry: ConfigEntry,
    ) -> None:
        """Initialize a generic RainMachine switch."""
        super().__init__(coordinator, controller, type(self).__name__)

        self._attr_is_on = False
        self._attr_name = name
        self._data = coordinator.data[uid]
        self._entry = entry
        self._is_active = True
        self._uid = uid

    @property
    def available(self) -> bool:
        """Return True if entity is available."""
        return super().available and self._is_active

    @property
    def unique_id(self) -> str:
        """Return a unique, Home Assistant friendly identifier for this entity."""
        return f"{super().unique_id}_{self._uid}"

    async def _async_run_switch_coroutine(self, api_coro: Coroutine) -> None:
        """Run a coroutine to toggle the switch."""
        try:
            resp = await api_coro
        except RequestError as err:
            LOGGER.error(
                'Error while toggling %s "%s": %s',
                self._entity_type,
                self.unique_id,
                err,
            )
            return

        if resp["statusCode"] != 0:
            LOGGER.error(
                'Error while toggling %s "%s": %s',
                self._entity_type,
                self.unique_id,
                resp["message"],
            )
            return

        # Because of how inextricably linked programs and zones are, anytime one is
        # toggled, we make sure to update the data of both coordinators:
        self.hass.async_create_task(
            async_update_programs_and_zones(self.hass, self._entry)
        )

    async def async_disable_program(self, *, program_id: int) -> None:
        """Disable a program."""
        await self._controller.programs.disable(program_id)
        await async_update_programs_and_zones(self.hass, self._entry)

    async def async_disable_zone(self, *, zone_id: int) -> None:
        """Disable a zone."""
        await self._controller.zones.disable(zone_id)
        await async_update_programs_and_zones(self.hass, self._entry)

    async def async_enable_program(self, *, program_id: int) -> None:
        """Enable a program."""
        await self._controller.programs.enable(program_id)
        await async_update_programs_and_zones(self.hass, self._entry)

    async def async_enable_zone(self, *, zone_id: int) -> None:
        """Enable a zone."""
        await self._controller.zones.enable(zone_id)
        await async_update_programs_and_zones(self.hass, self._entry)

    async def async_pause_watering(self, *, seconds: int) -> None:
        """Pause watering for a set number of seconds."""
        await self._controller.watering.pause_all(seconds)
        await async_update_programs_and_zones(self.hass, self._entry)

    async def async_start_program(self, *, program_id: int) -> None:
        """Start a particular program."""
        await self._controller.programs.start(program_id)
        await async_update_programs_and_zones(self.hass, self._entry)

    async def async_start_zone(self, *, zone_id: int, zone_run_time: int) -> None:
        """Start a particular zone for a certain amount of time."""
        await self._controller.zones.start(zone_id, zone_run_time)
        await async_update_programs_and_zones(self.hass, self._entry)

    async def async_stop_all(self) -> None:
        """Stop all watering."""
        await self._controller.watering.stop_all()
        await async_update_programs_and_zones(self.hass, self._entry)

    async def async_stop_program(self, *, program_id: int) -> None:
        """Stop a program."""
        await self._controller.programs.stop(program_id)
        await async_update_programs_and_zones(self.hass, self._entry)

    async def async_stop_zone(self, *, zone_id: int) -> None:
        """Stop a zone."""
        await self._controller.zones.stop(zone_id)
        await async_update_programs_and_zones(self.hass, self._entry)

    async def async_unpause_watering(self) -> None:
        """Unpause watering."""
        await self._controller.watering.unpause_all()
        await async_update_programs_and_zones(self.hass, self._entry)

    @callback
    def update_from_latest_data(self) -> None:
        """Update the state."""
        self._data = self.coordinator.data[self._uid]
        self._is_active = self._data["active"]


class RainMachineProgram(RainMachineSwitch):
    """A RainMachine program."""

    @property
    def zones(self) -> list:
        """Return a list of active zones associated with this program."""
        return [z for z in self._data["wateringTimes"] if z["active"]]

    async def async_turn_off(self, **kwargs: dict[str, Any]) -> None:
        """Turn the program off."""
        await self._async_run_switch_coroutine(
            self._controller.programs.stop(self._uid)
        )

    async def async_turn_on(self, **kwargs: dict[str, Any]) -> None:
        """Turn the program on."""
        await self._async_run_switch_coroutine(
            self._controller.programs.start(self._uid)
        )

    @callback
    def update_from_latest_data(self) -> None:
        """Update the state."""
        super().update_from_latest_data()

        self._attr_is_on = bool(self._data["status"])

        next_run: str | None = None
        if self._data.get("nextRun") is not None:
            next_run = datetime.strptime(
                f"{self._data['nextRun']} {self._data['startTime']}",
                "%Y-%m-%d %H:%M",
            ).isoformat()

        self._attr_extra_state_attributes.update(
            {
                ATTR_ID: self._uid,
                ATTR_NEXT_RUN: next_run,
                ATTR_SOAK: self.coordinator.data[self._uid].get("soak"),
                ATTR_STATUS: RUN_STATUS_MAP[self.coordinator.data[self._uid]["status"]],
                ATTR_ZONES: ", ".join(z["name"] for z in self.zones),
            }
        )


class RainMachineZone(RainMachineSwitch):
    """A RainMachine zone."""

    async def async_turn_off(self, **kwargs: dict[str, Any]) -> None:
        """Turn the zone off."""
        await self._async_run_switch_coroutine(self._controller.zones.stop(self._uid))

    async def async_turn_on(self, **kwargs: dict[str, Any]) -> None:
        """Turn the zone on."""
        await self._async_run_switch_coroutine(
            self._controller.zones.start(
                self._uid,
                self._entry.options[CONF_ZONE_RUN_TIME],
            )
        )

    @callback
    def update_from_latest_data(self) -> None:
        """Update the state."""
        super().update_from_latest_data()

        self._attr_is_on = bool(self._data["state"])

        self._attr_extra_state_attributes.update(
            {
                ATTR_STATUS: RUN_STATUS_MAP[self._data["state"]],
                ATTR_AREA: self._data.get("waterSense").get("area"),
                ATTR_CURRENT_CYCLE: self._data.get("cycle"),
                ATTR_FIELD_CAPACITY: self._data.get("waterSense").get("fieldCapacity"),
                ATTR_ID: self._data["uid"],
                ATTR_NO_CYCLES: self._data.get("noOfCycles"),
                ATTR_PRECIP_RATE: self._data.get("waterSense").get("precipitationRate"),
                ATTR_RESTRICTIONS: self._data.get("restriction"),
                ATTR_SLOPE: SLOPE_TYPE_MAP.get(self._data.get("slope")),
                ATTR_SOIL_TYPE: SOIL_TYPE_MAP.get(self._data.get("soil")),
                ATTR_SPRINKLER_TYPE: SPRINKLER_TYPE_MAP.get(self._data.get("group_id")),
                ATTR_SUN_EXPOSURE: SUN_EXPOSURE_MAP.get(self._data.get("sun")),
                ATTR_TIME_REMAINING: self._data.get("remaining"),
                ATTR_VEGETATION_TYPE: VEGETATION_MAP.get(self._data.get("type")),
            }
        )