ha-babytracker/custom_components/babytracker/services.py

"""Actions that add entries to a baby's log."""

from __future__ import annotations

from collections.abc import Callable
from datetime import datetime, timedelta
from typing import Any, Mapping

import requests

from homeassistant.config_entries import ConfigEntryState
from homeassistant.core import HomeAssistant, ServiceCall, ServiceResponse, SupportsResponse, callback
from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
from homeassistant.helpers import config_validation as cv, device_registry as dr
from homeassistant.util import dt as dt_util
from homeassistant.util.unit_system import US_CUSTOMARY_SYSTEM

from .babytracker_api import (
    Baby,
    BabyTrackerError,
    DiaperAmount,
    DiaperStatus,
    NursingFinishSide,
    bath,
    diaper,
    formula,
    growth,
    journal,
    joy,
    nursing,
    pumped,
    sleep,
    temperature,
)
from .compat import vol
from .const import DOMAIN
from .coordinator import BabyTrackerCoordinator

_WHO_AND_WHEN = {
    vol.Optional("baby"): cv.string,
    vol.Optional("time"): cv.datetime,
}
_BASE = {**_WHO_AND_WHEN, vol.Optional("note", default=""): cv.string}
_POSITIVE_FLOAT = vol.All(vol.Coerce(float), vol.Range(min=0))
_MINUTES = vol.All(vol.Coerce(int), vol.Range(min=0))


def _names(enum: Any, *excluded: str) -> list[str]:
    return [member.name.lower() for member in enum if member.name.lower() not in excluded]


def _any_of(*keys: str) -> Callable[[dict], dict]:
    def validate(data: dict) -> dict:
        if not any(data.get(key) for key in keys):
            raise vol.Invalid(f"at least one of {', '.join(keys)} is required")
        return data

    return validate


def _english(hass: HomeAssistant, unit: str | None, english_unit: str) -> bool:
    """Use the given unit, or else Home Assistant's unit system."""
    if unit is None:
        return hass.config.units is US_CUSTOMARY_SYSTEM
    return unit == english_unit


def _build_diaper(hass: HomeAssistant, baby: Mapping, data: Mapping, common: dict) -> dict:
    amount = DiaperAmount[data["amount"].upper()] if "amount" in data else DiaperAmount.NOSET
    return diaper(baby, DiaperStatus[data["status"].upper()], amount, **common)


def _build_sleep(hass: HomeAssistant, baby: Mapping, data: Mapping, common: dict) -> dict:
    return sleep(baby, data["duration"], **common)


def _build_bottle(hass: HomeAssistant, baby: Mapping, data: Mapping, common: dict) -> dict:
    build = pumped if data["kind"] == "pumped" else formula
    return build(baby, data["amount"], english=_english(hass, data.get("unit"), "oz"), **common)


def _build_nursing(hass: HomeAssistant, baby: Mapping, data: Mapping, common: dict) -> dict:
    return nursing(
        baby,
        data.get("left_minutes", 0),
        data.get("right_minutes", 0),
        data.get("both_minutes", 0),
        NursingFinishSide[data["finish_side"].upper()] if "finish_side" in data else NursingFinishSide.NOT_SET,
        **common,
    )


def _build_temperature(hass: HomeAssistant, baby: Mapping, data: Mapping, common: dict) -> dict:
    return temperature(baby, data["value"], english=_english(hass, data.get("unit"), "F"), **common)


def _build_growth(hass: HomeAssistant, baby: Mapping, data: Mapping, common: dict) -> dict:
    return growth(
        baby,
        data.get("weight"),
        data.get("length"),
        data.get("head"),
        english=_english(hass, data.get("units"), "english"),
        **common,
    )


def _build_bath(hass: HomeAssistant, baby: Mapping, data: Mapping, common: dict) -> dict:
    return bath(baby, **common)


def _build_joy(hass: HomeAssistant, baby: Mapping, data: Mapping, common: dict) -> dict:
    return joy(baby, **common)


def _build_journal(hass: HomeAssistant, baby: Mapping, data: Mapping, common: dict) -> dict:
    return journal(baby, **common)


def _elapsed(data: Mapping) -> int:
    return sum(data.get(key, 0) for key in ("duration", "left_minutes", "right_minutes", "both_minutes"))


_SERVICES: dict[str, tuple[Callable, Any]] = {
    "log_diaper": (
        _build_diaper,
        vol.Schema(
            {
                **_BASE,
                vol.Required("status"): vol.In(_names(DiaperStatus)),
                vol.Optional("amount"): vol.In(_names(DiaperAmount, "none", "noset")),
            }
        ),
    ),
    "log_sleep": (_build_sleep, vol.Schema({**_BASE, vol.Required("duration"): _MINUTES})),
    "log_bottle": (
        _build_bottle,
        vol.Schema(
            {
                **_BASE,
                vol.Required("amount"): _POSITIVE_FLOAT,
                vol.Optional("kind", default="formula"): vol.In(["formula", "pumped"]),
                vol.Optional("unit"): vol.In(["ml", "oz"]),
            }
        ),
    ),
    "log_nursing": (
        _build_nursing,
        vol.All(
            vol.Schema(
                {
                    **_BASE,
                    vol.Optional("left_minutes"): _MINUTES,
                    vol.Optional("right_minutes"): _MINUTES,
                    vol.Optional("both_minutes"): _MINUTES,
                    vol.Optional("finish_side"): vol.In(["left", "right"]),
                }
            ),
            _any_of("left_minutes", "right_minutes", "both_minutes"),
        ),
    ),
    "log_temperature": (
        _build_temperature,
        vol.Schema({**_BASE, vol.Required("value"): vol.Coerce(float), vol.Optional("unit"): vol.In(["C", "F"])}),
    ),
    "log_growth": (
        _build_growth,
        vol.All(
            vol.Schema(
                {
                    **_BASE,
                    vol.Optional("weight"): _POSITIVE_FLOAT,
                    vol.Optional("length"): _POSITIVE_FLOAT,
                    vol.Optional("head"): _POSITIVE_FLOAT,
                    vol.Optional("units"): vol.In(["metric", "english"]),
                }
            ),
            _any_of("weight", "length", "head"),
        ),
    ),
    "log_bath": (_build_bath, vol.Schema(_BASE)),
    "log_joy": (_build_joy, vol.Schema(_BASE)),
    "log_journal": (
        _build_journal,
        vol.Schema({**_WHO_AND_WHEN, vol.Required("note"): vol.All(cv.string, vol.Length(min=1))}),
    ),
}


def _coordinators(hass: HomeAssistant) -> list[BabyTrackerCoordinator]:
    return [
        entry.runtime_data
        for entry in hass.config_entries.async_entries(DOMAIN)
        if entry.state is ConfigEntryState.LOADED
    ]


def _resolve_baby(hass: HomeAssistant, wanted: str | None) -> tuple[BabyTrackerCoordinator, Baby]:
    """Find a baby by device ID or name, or the only baby if none was named."""
    known = [(coordinator, baby) for coordinator in _coordinators(hass) for baby in coordinator.data.babies.values()]
    if wanted is None:
        if len(known) == 1:
            return known[0]
        raise ServiceValidationError(translation_domain=DOMAIN, translation_key="baby_required")

    baby_id = None
    if (device := dr.async_get(hass).async_get(wanted)) is not None:
        baby_id = next((key for domain, key in device.identifiers if domain == DOMAIN), None)
    for coordinator, baby in known:
        if baby.id == baby_id or baby.name.casefold() == wanted.casefold():
            return coordinator, baby
    raise ServiceValidationError(
        translation_domain=DOMAIN, translation_key="baby_not_found", translation_placeholders={"baby": wanted}
    )


def _common_fields(data: Mapping) -> dict[str, Any]:
    """Fields every record shares. Without a time, the activity has just finished."""
    moment: datetime = (
        dt_util.as_utc(data["time"]) if "time" in data else dt_util.utcnow() - timedelta(minutes=_elapsed(data))
    )
    local = dt_util.as_local(moment)
    offset = (local.utcoffset() or timedelta(0)) - (local.dst() or timedelta(0))
    return {"time": moment, "note": data.get("note", ""), "timezone": int(offset.total_seconds())}


@callback
def async_setup_services(hass: HomeAssistant) -> None:
    def register(name: str, build: Callable, schema: Any) -> None:
        async def handle(call: ServiceCall) -> ServiceResponse:
            coordinator, baby = _resolve_baby(hass, call.data.get("baby"))
            record = build(hass, baby.record, call.data, _common_fields(call.data))
            try:
                record_id = await coordinator.async_add_record(record)
            except (BabyTrackerError, requests.RequestException) as err:
                raise HomeAssistantError(
                    translation_domain=DOMAIN,
                    translation_key="upload_failed",
                    translation_placeholders={"error": str(err)},
                ) from err
            return {"id": record_id} if call.return_response else None

        hass.services.async_register(DOMAIN, name, handle, schema=schema, supports_response=SupportsResponse.OPTIONAL)

    for name, (build, schema) in _SERVICES.items():
        register(name, build, schema)