"""Button platform for TaskMate integration.""" from __future__ import annotations import logging from homeassistant.components.button import ButtonEntity from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity import DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import DOMAIN from .coordinator import TaskMateCoordinator from .models import Child, Chore, Reward _LOGGER = logging.getLogger(__name__) async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up TaskMate buttons.""" coordinator: TaskMateCoordinator = hass.data[DOMAIN][entry.entry_id] entities: list[ButtonEntity] = [] # Add complete chore buttons for each child/chore combination children = coordinator.data.get("children", []) chores = coordinator.data.get("chores", []) rewards = coordinator.data.get("rewards", []) for child in children: # Chore completion buttons for chore in chores: if getattr(chore, "assignment_mode", "everyone") == "unassigned": continue if not chore.assigned_to or child.id in chore.assigned_to: entities.append( CompleteChoreButton(coordinator, entry, child, chore) ) # Reward claim buttons for reward in rewards: entities.append( ClaimRewardButton(coordinator, entry, child, reward) ) # Track which entity combos already exist tracked_combos: set[str] = set() for child in children: for chore in chores: if getattr(chore, "assignment_mode", "everyone") == "unassigned": continue if not chore.assigned_to or child.id in chore.assigned_to: tracked_combos.add(f"{child.id}_{chore.id}_complete") for reward in rewards: tracked_combos.add(f"{child.id}_{reward.id}_claim") async_add_entities(entities) # Set up listener to add buttons for new children/chores/rewards @callback def async_update_entities() -> None: """Add button entities for newly created children, chores, or rewards.""" new_entities: list[ButtonEntity] = [] current_children = coordinator.data.get("children", []) current_chores = coordinator.data.get("chores", []) current_rewards = coordinator.data.get("rewards", []) for child in current_children: for chore in current_chores: if getattr(chore, "assignment_mode", "everyone") == "unassigned": continue if not chore.assigned_to or child.id in chore.assigned_to: key = f"{child.id}_{chore.id}_complete" if key not in tracked_combos: new_entities.append( CompleteChoreButton(coordinator, entry, child, chore) ) tracked_combos.add(key) for reward in current_rewards: key = f"{child.id}_{reward.id}_claim" if key not in tracked_combos: new_entities.append( ClaimRewardButton(coordinator, entry, child, reward) ) tracked_combos.add(key) if new_entities: async_add_entities(new_entities) coordinator.async_add_listener(async_update_entities) class TaskMateBaseButton(CoordinatorEntity, ButtonEntity): """Base class for TaskMate buttons.""" def __init__( self, coordinator: TaskMateCoordinator, entry: ConfigEntry, ) -> None: """Initialize the button.""" super().__init__(coordinator) self._entry = entry @property def device_info(self) -> DeviceInfo: """Return device info.""" return DeviceInfo( identifiers={(DOMAIN, self._entry.entry_id)}, name="TaskMate", manufacturer="TaskMate", model="Family Chore Manager", ) class CompleteChoreButton(TaskMateBaseButton): """Button to mark a chore as completed.""" def __init__( self, coordinator: TaskMateCoordinator, entry: ConfigEntry, child: Child, chore: Chore, ) -> None: """Initialize the button.""" super().__init__(coordinator, entry) self.child_id = child.id self.chore_id = chore.id self._attr_unique_id = f"{entry.entry_id}_{child.id}_{chore.id}_complete" self._attr_name = f"{child.name}: Complete {chore.name}" @property def icon(self) -> str: """Return the icon.""" chore = self.coordinator.get_chore(self.chore_id) # Chores gained an optional icon in #683, defaulting to "". Fall back on # falsiness, not on the attribute being absent — otherwise every chore # without a picture gets a blank button icon. return (getattr(chore, 'icon', "") or "mdi:check-circle") if chore else "mdi:check-circle" @property def extra_state_attributes(self) -> dict: """Return additional attributes.""" child = self.coordinator.get_child(self.child_id) chore = self.coordinator.get_chore(self.chore_id) if not child or not chore: return {} return { "child_id": child.id, "child_name": child.name, "chore_id": chore.id, "chore_name": chore.name, "points": chore.points, "requires_approval": chore.requires_approval, } async def async_press(self) -> None: """Handle the button press.""" try: await self.coordinator.async_complete_chore(self.chore_id, self.child_id) except ValueError as err: _LOGGER.error("Failed to complete chore: %s", err) class ClaimRewardButton(TaskMateBaseButton): """Button to claim a reward.""" def __init__( self, coordinator: TaskMateCoordinator, entry: ConfigEntry, child: Child, reward: Reward, ) -> None: """Initialize the button.""" super().__init__(coordinator, entry) self.child_id = child.id self.reward_id = reward.id self._attr_unique_id = f"{entry.entry_id}_{child.id}_{reward.id}_claim" self._attr_name = f"{child.name}: Claim {reward.name}" @property def icon(self) -> str: """Return the icon.""" reward = self.coordinator.get_reward(self.reward_id) return reward.icon if reward else "mdi:gift" def _get_committed_points(self, child_id: str) -> int: """Calculate points committed by pending reward claims for a child. Pool-mode pending claims are skipped because their cost was already deducted from child.points at allocation time. """ pending_claims = self.coordinator.data.get("pending_reward_claims", []) committed = 0 for c in pending_claims: if c.child_id == child_id and not self.coordinator.is_pool_mode_claim(c): reward = self.coordinator.get_reward(c.reward_id) if reward: committed += reward.cost return committed @property def available(self) -> bool: """Return if button is available (accounting for committed points).""" child = self.coordinator.get_child(self.child_id) reward = self.coordinator.get_reward(self.reward_id) if not child or not reward: return False available_points = child.points - self._get_committed_points(child.id) return available_points >= reward.cost @property def extra_state_attributes(self) -> dict: """Return additional attributes.""" child = self.coordinator.get_child(self.child_id) reward = self.coordinator.get_reward(self.reward_id) if not child or not reward: return {} committed = self._get_committed_points(child.id) available_points = child.points - committed can_afford = available_points >= reward.cost return { "child_id": child.id, "child_name": child.name, "reward_id": reward.id, "reward_name": reward.name, "cost": reward.cost, "child_points": child.points, "committed_points": committed, "available_points": available_points, "can_afford": can_afford, "points_needed": max(0, reward.cost - available_points), } async def async_press(self) -> None: """Handle the button press.""" try: await self.coordinator.async_claim_reward(self.reward_id, self.child_id) except ValueError as err: _LOGGER.error("Failed to claim reward: %s", err)