MedMate-Medicine-Tracker/sensor.py

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

import logging
from datetime import datetime
from typing import Any

from homeassistant.components.sensor import (
    SensorEntity,
    SensorEntityDescription,
    SensorStateClass,
    SensorDeviceClass,
)
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_ACTIVE_INGREDIENT,
    CONF_STRENGTH,
    CONF_PACK_SIZE,
    CONF_SCHEDULE,
    CONF_PRESCRIPTION,
    CONF_ISSUE_DATE,
    CONF_EXPIRY_DATE,
    CONF_DOCTOR,
    CONF_TOTAL_REPEATS,
    CONF_REPEATS_LEFT,
    ATTR_ACTIVE_INGREDIENT,
    ATTR_STRENGTH,
    ATTR_PACK_SIZE,
    ATTR_INVENTORY,
    ATTR_SCHEDULE,
    ATTR_PRESCRIPTION,
    ATTR_ISSUE_DATE,
    ATTR_EXPIRY_DATE,
    ATTR_DOCTOR,
    ATTR_TOTAL_REPEATS,
    ATTR_REPEATS_LEFT,
    ATTR_PRESCRIPTION_ACTIVE,
    ATTR_NEXT_DOSE,
    ATTR_LAST_TAKEN,
)
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 sensors."""
    coordinator = hass.data[DOMAIN][config_entry.entry_id]
    
    sensors = [
        MedMateInventorySensor(coordinator, config_entry),
        MedMateRepeatsSensor(coordinator, config_entry),
        MedMateNextDoseSensor(coordinator, config_entry),
    ]
    
    async_add_entities(sensors)


class MedMateBaseSensor(CoordinatorEntity, SensorEntity):
    """Base sensor for MedMate."""

    def __init__(
        self,
        coordinator: MedMateDataUpdateCoordinator,
        config_entry: ConfigEntry,
        description: SensorEntityDescription,
    ) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator)
        self.entity_description = description
        self.config_entry = config_entry
        self._medicine_id = config_entry.data.get("medicine_id")
        
        # Set unique_id and entity_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 MedMateInventorySensor(MedMateBaseSensor):
    """Sensor for medicine inventory."""

    def __init__(
        self,
        coordinator: MedMateDataUpdateCoordinator,
        config_entry: ConfigEntry,
    ) -> None:
        """Initialize the inventory sensor."""
        description = SensorEntityDescription(
            key="inventory",
            name="Inventory",
            icon="mdi:pill",
            state_class=SensorStateClass.MEASUREMENT,
            native_unit_of_measurement="pills",
        )
        super().__init__(coordinator, config_entry, description)

    @property
    def native_value(self) -> int:
        """Return the inventory count."""
        return self.coordinator.data.get("inventory", 0)

    @property
    def extra_state_attributes(self) -> dict[str, Any]:
        """Return additional state attributes."""
        data = self.coordinator.data
        schedule = data.get(CONF_SCHEDULE, {})
        prescription = data.get(CONF_PRESCRIPTION, {})
        
        attrs = {
            ATTR_ACTIVE_INGREDIENT: data.get(CONF_ACTIVE_INGREDIENT),
            ATTR_STRENGTH: data.get(CONF_STRENGTH),
            ATTR_PACK_SIZE: data.get(CONF_PACK_SIZE),
            ATTR_SCHEDULE: schedule,
            ATTR_PRESCRIPTION_ACTIVE: data.get("prescription_active", False),
        }
        
        # Add prescription details if available
        if prescription:
            attrs.update({
                ATTR_ISSUE_DATE: prescription.get(CONF_ISSUE_DATE),
                ATTR_EXPIRY_DATE: prescription.get(CONF_EXPIRY_DATE),
                ATTR_DOCTOR: prescription.get(CONF_DOCTOR),
                ATTR_TOTAL_REPEATS: prescription.get(CONF_TOTAL_REPEATS),
                ATTR_REPEATS_LEFT: prescription.get(CONF_REPEATS_LEFT),
            })
            
        return attrs


class MedMateRepeatsSensor(MedMateBaseSensor):
    """Sensor for prescription repeats left."""

    def __init__(
        self,
        coordinator: MedMateDataUpdateCoordinator,
        config_entry: ConfigEntry,
    ) -> None:
        """Initialize the repeats sensor."""
        description = SensorEntityDescription(
            key="repeats_left",
            name="Repeats Left",
            icon="mdi:repeat",
            state_class=SensorStateClass.MEASUREMENT,
            native_unit_of_measurement="repeats",
        )
        super().__init__(coordinator, config_entry, description)

    @property
    def native_value(self) -> int:
        """Return the number of repeats left."""
        prescription = self.coordinator.data.get(CONF_PRESCRIPTION, {})
        return prescription.get(CONF_REPEATS_LEFT, 0)

    @property
    def extra_state_attributes(self) -> dict[str, Any]:
        """Return additional state attributes."""
        data = self.coordinator.data
        prescription = data.get(CONF_PRESCRIPTION, {})
        
        return {
            ATTR_TOTAL_REPEATS: prescription.get(CONF_TOTAL_REPEATS, 0),
            ATTR_PRESCRIPTION_ACTIVE: data.get("prescription_active", False),
            ATTR_ISSUE_DATE: prescription.get(CONF_ISSUE_DATE),
            ATTR_EXPIRY_DATE: prescription.get(CONF_EXPIRY_DATE),
            ATTR_DOCTOR: prescription.get(CONF_DOCTOR),
        }


class MedMateNextDoseSensor(MedMateBaseSensor):
    """Sensor for next dose time."""

    def __init__(
        self,
        coordinator: MedMateDataUpdateCoordinator,
        config_entry: ConfigEntry,
    ) -> None:
        """Initialize the next dose sensor."""
        description = SensorEntityDescription(
            key="next_dose",
            name="Next Dose",
            icon="mdi:clock-outline",
            device_class=SensorDeviceClass.TIMESTAMP,
        )
        super().__init__(coordinator, config_entry, description)

    @property
    def native_value(self) -> datetime | None:
        """Return the next dose time."""
        next_dose = self.coordinator.data.get("next_dose")
        if next_dose and isinstance(next_dose, str):
            try:
                return datetime.fromisoformat(next_dose)
            except ValueError:
                return None
        return next_dose

    @property
    def extra_state_attributes(self) -> dict[str, Any]:
        """Return additional state attributes."""
        data = self.coordinator.data
        schedule = data.get(CONF_SCHEDULE, {})
        last_taken = data.get("last_taken", {})
        
        # Get most recent dose taken
        most_recent_dose = None
        if last_taken:
            recent_doses = sorted(last_taken.values(), reverse=True)
            if recent_doses:
                try:
                    most_recent_dose = datetime.fromisoformat(recent_doses[0])
                except (ValueError, TypeError):
                    most_recent_dose = None
        
        return {
            ATTR_SCHEDULE: schedule,
            ATTR_LAST_TAKEN: most_recent_dose,
            "dose_due": data.get("dose_due", False),
        }