MedMate-Medicine-Tracker/button.py

"""Button platform for MedMate integration."""
from __future__ import annotations

import logging
from typing import Any

from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import CoordinatorEntity

from .const import (
    DOMAIN,
    CONF_MEDICINE_NAME,
    CONF_SCHEDULE,
    CONF_TIMES,
    TIME_SLOTS,
)
from .coordinator import MedMateDataUpdateCoordinator

_LOGGER = logging.getLogger(__name__)


async def async_setup_entry(
    hass: HomeAssistant,
    config_entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up MedMate buttons."""
    coordinator = hass.data[DOMAIN][config_entry.entry_id]
    
    # Wait for first update to ensure we have clean data
    if not coordinator.data:
        await coordinator.async_request_refresh()
    
    buttons = [
        MedMateFillPrescriptionButton(coordinator, config_entry),
    ]
    
    # Safely get scheduled times with error handling
    try:
        schedule = coordinator.data.get(CONF_SCHEDULE, {})
        scheduled_times = schedule.get(CONF_TIMES, []) if isinstance(schedule, dict) else []
        
        # Validate that scheduled_times is a list and contains expected values
        if isinstance(scheduled_times, list):
            for time_slot in scheduled_times:
                if isinstance(time_slot, str) and time_slot in TIME_SLOTS:
                    buttons.append(
                        MedMateTakeDoseButton(coordinator, config_entry, time_slot)
                    )
        else:
            _LOGGER.warning("Invalid scheduled_times format: %s", type(scheduled_times))
            
    except Exception as err:
        _LOGGER.error("Error setting up dose buttons: %s", err)
        # Continue with just the fill prescription button
    
    async_add_entities(buttons)


class MedMateBaseButton(CoordinatorEntity, ButtonEntity):
    """Base button for MedMate."""

    def __init__(
        self,
        coordinator: MedMateDataUpdateCoordinator,
        config_entry: ConfigEntry,
        description: ButtonEntityDescription,
    ) -> None:
        """Initialize the button."""
        super().__init__(coordinator)
        self.entity_description = description
        self.config_entry = config_entry
        self._medicine_id = config_entry.data.get("medicine_id")
        
        # Set unique_id
        self._attr_unique_id = f"{self._medicine_id}_{description.key}"
        
    @property
    def device_info(self) -> dict[str, Any]:
        """Return device information."""
        medicine_name = self.coordinator.data.get(CONF_MEDICINE_NAME, "Unknown Medicine")
        return {
            "identifiers": {(DOMAIN, self._medicine_id)},
            "name": f"MedMate - {medicine_name}",
            "manufacturer": "MedMate",
            "model": "Medicine Tracker",
            "sw_version": "1.0.0",
        }

    @property
    def available(self) -> bool:
        """Return if entity is available."""
        return self.coordinator.last_update_success and bool(self.coordinator.data)


class MedMateFillPrescriptionButton(MedMateBaseButton):
    """Button to fill prescription."""

    def __init__(
        self,
        coordinator: MedMateDataUpdateCoordinator,
        config_entry: ConfigEntry,
    ) -> None:
        """Initialize the fill prescription button."""
        description = ButtonEntityDescription(
            key="fill_prescription",
            name="Fill Prescription",
            icon="mdi:pill",
        )
        super().__init__(coordinator, config_entry, description)

    async def async_press(self) -> None:
        """Handle the button press."""
        success = await self.coordinator.async_fill_prescription()
        
        if success:
            _LOGGER.info(
                "Prescription filled for medicine: %s",
                self.coordinator.data.get(CONF_MEDICINE_NAME)
            )
        else:
            _LOGGER.warning(
                "Failed to fill prescription for medicine: %s",
                self.coordinator.data.get(CONF_MEDICINE_NAME)
            )

    @property
    def available(self) -> bool:
        """Return if button is available."""
        if not super().available:
            return False
            
        # Only available if prescription is active and has repeats left
        return self.coordinator.data.get("prescription_active", False)


class MedMateTakeDoseButton(MedMateBaseButton):
    """Button to take a dose."""

    def __init__(
        self,
        coordinator: MedMateDataUpdateCoordinator,
        config_entry: ConfigEntry,
        time_slot: str,
    ) -> None:
        """Initialize the take dose button."""
        self.time_slot = time_slot
        time_slot_name = TIME_SLOTS.get(time_slot, time_slot.title())
        
        description = ButtonEntityDescription(
            key=f"take_dose_{time_slot}",
            name=f"Take {time_slot_name} Dose",
            icon="mdi:hand-heart",
        )
        super().__init__(coordinator, config_entry, description)

    async def async_press(self) -> None:
        """Handle the button press."""
        success = await self.coordinator.async_take_dose(self.time_slot)
        
        if success:
            _LOGGER.info(
                "Dose taken for medicine: %s at %s",
                self.coordinator.data.get(CONF_MEDICINE_NAME),
                self.time_slot
            )
        else:
            _LOGGER.warning(
                "Failed to record dose for medicine: %s at %s",
                self.coordinator.data.get(CONF_MEDICINE_NAME),
                self.time_slot
            )

    @property
    def available(self) -> bool:
        """Return if button is available."""
        if not super().available:
            return False
            
        # Only available if there's inventory
        inventory = self.coordinator.data.get("inventory", 0)
        return inventory > 0

    @property
    def extra_state_attributes(self) -> dict[str, Any]:
        """Return additional state attributes."""
        return {
            "time_slot": self.time_slot,
            "current_inventory": self.coordinator.data.get("inventory", 0),
        }